feat: introduce enterprise to Client constructor and GOOGLE_GENAI_USE_ENTERPRISE

PiperOrigin-RevId: 899811032

Author: yyyu-google

README.md

@@ -12,7 +12,7 @@
 Google Gen AI Python SDK provides an interface for developers to integrate
 Google's generative models into their Python applications. It supports the
 [Gemini Developer API](https://ai.google.dev/gemini-api/docs) and
-[Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/overview)
+[Gemini Enterprise Agent Platform](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/overview)
 APIs.
 
 ## Code Generation

google/genai/_api_client.py

@@ -579,11 +579,32 @@ def __init__(
     self.vertexai = vertexai
     self.custom_base_url = None
     if self.vertexai is None:
-      if os.environ.get('GOOGLE_GENAI_USE_VERTEXAI', '0').lower() in [
-          'true',
-          '1',
-      ]:
-        self.vertexai = True
+      env_enterprise_str = os.environ.get('GOOGLE_GENAI_USE_ENTERPRISE', None)
+      env_vertexai_str = os.environ.get('GOOGLE_GENAI_USE_VERTEXAI', None)
+
+      env_enterprise = None
+      if env_enterprise_str is not None:
+        env_enterprise = env_enterprise_str.lower() in ['true', '1']
+
+      env_vertexai = None
+      if env_vertexai_str is not None:
+        env_vertexai = env_vertexai_str.lower() in ['true', '1']
+
+      if (
+          env_enterprise is not None
+          and env_vertexai is not None
+          and env_enterprise != env_vertexai
+      ):
+        warnings.warn(
+            '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 env_enterprise is not None:
+        self.vertexai = env_enterprise
+      elif env_vertexai is not None:
+        self.vertexai = env_vertexai
 
     # Validate explicitly set initializer values.
     if (project or location) and api_key:

google/genai/client.py

@@ -313,37 +313,43 @@ class DebugConfig(pydantic.BaseModel):
 class Client:
   """Client for making synchronous requests.
 
-  Use this client to make a request to the Gemini Developer API or Vertex AI
-  API and then wait for the response.
+  Use this client to make a request to the Gemini Developer API or Gemini
+  Enterprise Agent Platform (previously Vertex AI API) and then wait for the
+  response.
 
   To initialize the client, provide the required arguments either directly
   or by using environment variables. Gemini API users and Vertex AI users in
-  express mode can provide API key by providing input argument
   `api_key="your-api-key"` or by defining `GOOGLE_API_KEY="your-api-key"` as an
   environment variable
 
-  Vertex AI API users can provide inputs argument as `vertexai=True,
+  Gemini Enterprise Agent Platform API users can provide inputs argument as
+  `enterprise=True,
   project="your-project-id", location="us-central1"` or by defining
-  `GOOGLE_GENAI_USE_VERTEXAI=true`, `GOOGLE_CLOUD_PROJECT` and
+  `GOOGLE_GENAI_USE_ENTERPRISE=true`, `GOOGLE_CLOUD_PROJECT` and
+  `GOOGLE_CLOUD_LOCATION` environment variables.
   `GOOGLE_CLOUD_LOCATION` environment variables.
 
   Attributes:
     api_key: The `API key <https://ai.google.dev/gemini-api/docs/api-key>`_ to
       use for authentication. Applies to the Gemini Developer API only.
-    vertexai: Indicates whether the client should use the Vertex AI API
-      endpoints. Defaults to False (uses Gemini Developer API endpoints).
-      Applies to the Vertex AI API only.
+    enterprise (bool): Indicates whether the client should use the Gemini
+      Enterprise Agent Platform endpoints (previously Vertex AI API).
+      Defaults to False (uses Gemini Developer API endpoints). When
+      `enterprise` and `vertexai` are both set, and they have conflicting
+      values, a `ValueError` will be raised.
+    vertexai (bool): Legacy flag for `enterprise`.
     credentials: The credentials to use for authentication when calling the
-      Vertex AI APIs. Credentials can be obtained from environment variables and
-      default credentials. For more information, see `Set up Application Default
-      Credentials
+      Gemini Enterprise Agent Platform APIs. Credentials can be obtained from
+      environment variables and default credentials. For more information, see
+      `Set up Application Default Credentials
       <https://cloud.google.com/docs/authentication/provide-credentials-adc>`_.
       Applies to the Vertex AI API only.
     project: The `Google Cloud project ID
       <https://cloud.google.com/vertex-ai/docs/start/cloud-environment>`_ to use
       for quota. Can be obtained from environment variables (for example,
       ``GOOGLE_CLOUD_PROJECT``). Applies to the Vertex AI API only.
-      Find your `Google Cloud project ID <https://cloud.google.com/resource-manager/docs/creating-managing-projects#identifying_projects>`_.
+      Find your `Google Cloud project ID
+      <https://cloud.google.com/resource-manager/docs/creating-managing-projects#identifying_projects>`_.
     location: The `location
       <https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations>`_
       to send API requests to (for example, ``us-central1``). Can be obtained
@@ -362,20 +368,21 @@ class Client:
 
     client = genai.Client(api_key='my-api-key')
 
-  Usage for the Vertex AI API:
+  Usage for the Gemini Enterprise Agent Platform API:
 
   .. code-block:: python
 
     from google import genai
 
     client = genai.Client(
-        vertexai=True, project='my-project-id', location='us-central1'
+        enterprise=True, project='my-project-id', location='us-central1'
     )
   """
 
   def __init__(
       self,
       *,
+      enterprise: Optional[bool] = None,
       vertexai: Optional[bool] = None,
       api_key: Optional[str] = None,
       credentials: Optional[google.auth.credentials.Credentials] = None,
@@ -387,9 +394,12 @@ def __init__(
     """Initializes the client.
 
     Args:
-       vertexai (bool): Indicates whether the client should use the Vertex AI
-         API endpoints. Defaults to False (uses Gemini Developer API endpoints).
-         Applies to the Vertex AI API only.
+       enterprise (bool): Indicates whether the client should use the Gemini
+         Enterprise Agent Platform endpoints (previously Vertex AI API).
+         Defaults to False (uses Gemini Developer API endpoints). When
+         `enterprise` and `vertexai` are both set, and they have conflicting
+         values, a `ValueError` will be raised.
+       vertexai (bool): Legacy flag for `enterprise`.
        api_key (str): The `API key
          <https://ai.google.dev/gemini-api/docs/api-key>`_ to use for
          authentication. Applies to the Gemini Developer API only.
@@ -414,18 +424,27 @@ def __init__(
     """
 
     self._debug_config = debug_config or DebugConfig()
+
+    if enterprise is not None and vertexai is not None and enterprise != vertexai:
+      raise ValueError(
+          'enterprise and vertexai flags have conflicting values, please set'
+          ' enterprise value only.'
+      )
+
+    resolved_vertexai = enterprise if enterprise is not None else vertexai
+
     if isinstance(http_options, dict):
       http_options = HttpOptions(**http_options)
 
-    base_url = get_base_url(vertexai or False, http_options)
+    base_url = get_base_url(resolved_vertexai or False, http_options)
     if base_url:
       if http_options:
         http_options.base_url = base_url
       else:
         http_options = HttpOptions(base_url=base_url)
 
     self._api_client = self._get_api_client(
-        vertexai=vertexai,
+        vertexai=resolved_vertexai,
         api_key=api_key,
         credentials=credentials,
         project=project,

google/genai/live.py

@@ -207,7 +207,7 @@ async def send_client_content(
       from google.genai import types
       import os
 
-      if os.environ.get('GOOGLE_GENAI_USE_VERTEXAI'):
+      if os.environ.get('GOOGLE_GENAI_USE_ENTERPRISE'):
         MODEL_NAME = 'gemini-2.0-flash-live-preview-04-09'
       else:
         MODEL_NAME = 'gemini-live-2.5-flash-preview';
@@ -279,7 +279,7 @@ async def send_realtime_input(
 
       import os
 
-      if os.environ.get('GOOGLE_GENAI_USE_VERTEXAI'):
+      if os.environ.get('GOOGLE_GENAI_USE_ENTERPRISE'):
         MODEL_NAME = 'gemini-2.0-flash-live-preview-04-09'
       else:
         MODEL_NAME = 'gemini-live-2.5-flash-preview';
@@ -374,7 +374,7 @@ async def send_tool_response(
 
       import os
 
-      if os.environ.get('GOOGLE_GENAI_USE_VERTEXAI'):
+      if os.environ.get('GOOGLE_GENAI_USE_ENTERPRISE'):
         MODEL_NAME = 'gemini-2.0-flash-live-preview-04-09'
       else:
         MODEL_NAME = 'gemini-live-2.5-flash-preview';

google/genai/tests/client/test_client_initialization.py

@@ -20,7 +20,6 @@
 import concurrent.futures
 import logging
 import os
-import requests
 import ssl
 from unittest import mock
 
@@ -29,14 +28,17 @@
 from google.auth import credentials
 import httpx
 import pytest
+import requests
 
 from ... import _api_client as api_client
 from ... import _base_url as base_url
 from ... import _replay_api_client as replay_api_client
 from ... import Client
 from ... import types
+
 try:
   import aiohttp
+
   AIOHTTP_NOT_INSTALLED = False
 except ImportError:
   AIOHTTP_NOT_INSTALLED = True
@@ -331,6 +333,84 @@ def test_vertexai_from_env_true(monkeypatch):
   assert client.models._api_client.location == location
 
 
+def test_enterprise_constructor_true():
+  client = Client(
+      enterprise=True, project="fake_project_id", location="fake-location"
+  )
+  assert client.models._api_client.vertexai
+
+
+def test_enterprise_constructor_false():
+  client = Client(enterprise=False, api_key="fake_api_key")
+  assert not client.models._api_client.vertexai
+
+
+def test_enterprise_constructor_conflict():
+  with pytest.raises(
+      ValueError,
+      match=(
+          "enterprise and vertexai flags have conflicting values, please set"
+          " enterprise value only."
+      ),
+  ):
+    Client(enterprise=True, vertexai=False)
+
+
+def test_enterprise_env_true(monkeypatch):
+  monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "true")
+  client = Client(project="fake_project_id", location="fake-location")
+  assert client.models._api_client.vertexai
+
+
+def test_enterprise_env_false(monkeypatch):
+  monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "false")
+  client = Client(api_key="fake_api_key")
+  assert not client.models._api_client.vertexai
+
+
+def test_enterprise_env_conflict_warning(monkeypatch):
+  monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "true")
+  monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "false")
+
+  with pytest.warns(
+      UserWarning,
+      match=(
+          "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."
+      ),
+  ):
+    # In BaseApiClient, resolving this warning.
+    client = Client(project="fake_project_id", location="fake-location")
+
+  assert client.models._api_client.vertexai
+
+
+def test_enterprise_constructor_precedence(monkeypatch):
+  monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "false")
+  client = Client(
+      enterprise=True, project="fake_project_id", location="fake-location"
+  )
+  assert client.models._api_client.vertexai
+
+
+def test_enterprise_precedence_over_vertexai_constructor():
+  client = Client(
+      enterprise=True,
+      vertexai=True,
+      project="fake_project_id",
+      location="fake-location",
+  )
+  assert client.models._api_client.vertexai
+
+
+def test_enterprise_env_precedence_over_vertexai_env(monkeypatch):
+  monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "false")
+  monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true")
+  client = Client(api_key="fake_api_key")
+  assert not client.models._api_client.vertexai
+
+
 def test_vertexai_from_constructor():
   project_id = "fake_project_id"
   location = "fake-location"
@@ -373,7 +453,9 @@ def mock_auth_default(scopes=None):
   monkeypatch.setattr(google.auth, "default", mock_auth_default)
   # Including a base_url override skips the check for having proj/location or
   # api_key set.
-  client = Client(vertexai=True, http_options={"base_url": "https://override.com/"})
+  client = Client(
+      vertexai=True, http_options={"base_url": "https://override.com/"}
+  )
   assert client.models._api_client.location is None
 
 
@@ -461,7 +543,7 @@ def test_vertexai_default_location_to_global_with_vertexai_base_url(
     m.setenv("GOOGLE_CLOUD_PROJECT", project_id)
     client = Client(
         vertexai=True,
-        http_options={'base_url': 'https://fake-url.googleapis.com'},
+        http_options={"base_url": "https://fake-url.googleapis.com"},
     )
     # Implicit project takes precedence over implicit api_key
     assert client.models._api_client.location == "global"
@@ -479,7 +561,7 @@ def test_vertexai_default_location_to_global_with_arbitrary_base_url(
     m.setenv("GOOGLE_CLOUD_PROJECT", project_id)
     client = Client(
         vertexai=True,
-        http_options={'base_url': 'https://fake-url.com'},
+        http_options={"base_url": "https://fake-url.com"},
     )
     # Implicit project takes precedence over implicit api_key
     assert not client.models._api_client.location
@@ -602,10 +684,7 @@ def test_vertexai_explicit_credentials(monkeypatch):
   monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "fake-location")
   monkeypatch.setenv("GOOGLE_API_KEY", "env_api_key")
 
-  client = Client(
-      vertexai=True,
-      credentials=creds
-  )
+  client = Client(vertexai=True, credentials=creds)
 
   assert client.models._api_client.vertexai
   assert client.models._api_client.project
@@ -1399,7 +1478,9 @@ def refresh_side_effect(request):
     )
     mock_request = mock.Mock(return_value=mock_http_response)
     monkeypatch.setattr(
-       google.auth.transport.requests.AuthorizedSession, "request", mock_request
+        google.auth.transport.requests.AuthorizedSession,
+        "request",
+        mock_request,
     )
   else:
     # Non-cloud environment w/o certificates uses httpx.Response