Introduce GEMINI_API_KEY env variable for API key, will take precedence over GOOGLE_API_KEY.

PiperOrigin-RevId: 737771755

Author: wanlin31

README.md

@@ -45,10 +45,12 @@ You can create a client by configuring the necessary environment variables.
 Configuration setup instructions depends on whether you're using the Gemini API
 on Vertex AI or the ML Dev Gemini API.
 
-**ML Dev Gemini API:** Set `GOOGLE_API_KEY` as shown below:
+**ML Dev Gemini API:** Set `GEMINI_API_KEY`(preferred) or `GOOGLE_API_KEY` as
+shown below:
 
 ```bash
-export GOOGLE_API_KEY='your-api-key'
+export GEMINI_API_KEY='your-api-key'
+
 ```
 
 **Vertex AI API:** Set `GOOGLE_GENAI_USE_VERTEXAI`, `GOOGLE_CLOUD_PROJECT`

google/genai/_api_client.py

@@ -287,7 +287,10 @@ def __init__(
     # Retrieve implicitly set values from the environment.
     env_project = os.environ.get('GOOGLE_CLOUD_PROJECT', None)
     env_location = os.environ.get('GOOGLE_CLOUD_LOCATION', None)
-    env_api_key = os.environ.get('GOOGLE_API_KEY', None)
+    # Gemini API key takes precedence over Google API key.
+    env_api_key = os.environ.get('GEMINI_API_KEY', None)
+    if not env_api_key or env_api_key == '':
+      env_api_key = os.environ.get('GOOGLE_API_KEY', None)
     self.project = project or env_project
     self.location = location or env_location
     self.api_key = api_key or env_api_key

google/genai/client.py

@@ -102,8 +102,9 @@ class Client:
   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
+  `api_key="your-api-key"` or via environment variables: `GOOGLE_API_KEY` or
+  `GEMINI_API_KEY`. When both environment variables are set `GEMINI_API_KEY`
+  takes precedence.
 
   Vertex AI API users can provide inputs argument as `vertexai=True,
   project="your-project-id", location="us-central1"` or by defining

google/genai/tests/client/test_client_initialization.py

@@ -26,16 +26,41 @@
 from ... import Client
 
 
-def test_ml_dev_from_env(monkeypatch):
+def test_ml_dev_from_google_api_key_env(monkeypatch):
   api_key = "google_api_key"
   monkeypatch.setenv("GOOGLE_API_KEY", api_key)
+  monkeypatch.setenv("GEMINI_API_KEY", "")
 
   client = Client()
 
   assert not client.models._api_client.vertexai
   assert client.models._api_client.api_key == api_key
   assert isinstance(client.models._api_client, api_client.BaseApiClient)
 
+def test_ml_dev_from_gemini_api_key_env(monkeypatch):
+  api_key = "gemini_api_key"
+  monkeypatch.setenv("GOOGLE_API_KEY", "")
+  monkeypatch.setenv("GEMINI_API_KEY", api_key)
+
+  client = Client()
+
+  assert not client.models._api_client.vertexai
+  assert client.models._api_client.api_key == api_key
+  assert isinstance(client.models._api_client, api_client.BaseApiClient)
+
+def test_ml_dev_from_env_precedence(monkeypatch):
+  # Gemini API key takes precedence over Google API key.
+  gemini_api_key = "gemini_api_key"
+  monkeypatch.setenv("GEMINI_API_KEY", gemini_api_key)
+  google_api_key = "google_api_key"
+  monkeypatch.setenv("GOOGLE_API_KEY", google_api_key)
+
+  client = Client()
+
+  assert not client.models._api_client.vertexai
+  assert client.models._api_client.api_key == gemini_api_key
+  assert isinstance(client.models._api_client, api_client.BaseApiClient)
+
 
 def test_ml_dev_from_constructor():
   api_key = "google_api_key"
@@ -271,6 +296,7 @@ def test_invalid_vertexai_constructor_empty(monkeypatch):
     monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "")
     monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "")
     monkeypatch.setenv("GOOGLE_API_KEY", "")
+    monkeypatch.setenv("GEMINI_API_KEY", "")
     def mock_auth_default(scopes=None):
       return None, None
 
@@ -281,6 +307,7 @@ def mock_auth_default(scopes=None):
 def test_invalid_mldev_constructor_empty(monkeypatch):
   with pytest.raises(ValueError):
     monkeypatch.setenv("GOOGLE_API_KEY", "")
+    monkeypatch.setenv("GEMINI_API_KEY", "")
     Client()
 
 
@@ -389,8 +416,8 @@ def test_mldev_explicit_arg_precedence(monkeypatch):
 
 
 def test_replay_client_ml_dev_from_env(monkeypatch, use_vertex: bool):
-  api_key = "google_api_key"
-  monkeypatch.setenv("GOOGLE_API_KEY", api_key)
+  gemini_api_key = "gemini_api_key"
+  monkeypatch.setenv("GEMINI_API_KEY", gemini_api_key)
   monkeypatch.setenv("GOOGLE_GENAI_CLIENT_MODE", "replay")
   api_type = "vertex" if use_vertex else "mldev"
   monkeypatch.setenv("GOOGLE_GENAI_REPLAY_ID", "test_replay_id." + api_type)
@@ -399,7 +426,7 @@ def test_replay_client_ml_dev_from_env(monkeypatch, use_vertex: bool):
   client = Client()
 
   assert not client.models._api_client.vertexai
-  assert client.models._api_client.api_key == api_key
+  assert client.models._api_client.api_key == gemini_api_key
   assert isinstance(
       client.models._api_client, replay_api_client.ReplayApiClient
   )
@@ -461,10 +488,11 @@ def test_vertexai_apikey_from_constructor(monkeypatch):
   assert isinstance(client.models._api_client, api_client.BaseApiClient)
 
 
-def test_vertexai_apikey_from_env(monkeypatch):
+def test_vertexai_apikey_from_google_api_key_env(monkeypatch):
   # Vertex AI Express mode uses API key on Vertex AI.
   api_key = "vertexai_api_key"
   monkeypatch.setenv("GOOGLE_API_KEY", api_key)
+  monkeypatch.setenv("GEMINI_API_KEY", "")
 
   # Due to proj/location taking precedence, need to clear proj/location env
   # variables.
@@ -480,6 +508,27 @@ def test_vertexai_apikey_from_env(monkeypatch):
   assert "aiplatform" in client._api_client._http_options["base_url"]
   assert isinstance(client.models._api_client, api_client.BaseApiClient)
 
+def test_vertexai_apikey_from_env_precedence(monkeypatch):
+  # Vertex AI Express mode uses gemnai API key on Vertex AI.
+  gemini_api_key = "gemini_api_key"
+  monkeypatch.setenv("GEMINI_API_KEY", gemini_api_key)
+  google_api_key = "google_api_key"
+  monkeypatch.setenv("GOOGLE_API_KEY", google_api_key)
+
+  # Due to proj/location taking precedence, need to clear proj/location env
+  # variables.
+  monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "")
+  monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "")
+
+  client = Client(vertexai=True)
+
+  assert client.models._api_client.vertexai
+  assert client.models._api_client.api_key == gemini_api_key
+  assert not client.models._api_client.project
+  assert not client.models._api_client.location
+  assert "aiplatform" in client._api_client._http_options["base_url"]
+  assert isinstance(client.models._api_client, api_client.BaseApiClient)
+
 
 def test_vertexai_apikey_invalid_constructor1():
   # Vertex AI Express mode uses API key on Vertex AI.
@@ -504,6 +553,7 @@ def test_vertexai_apikey_combo1(monkeypatch):
   monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", project_id)
   monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", location)
   monkeypatch.setenv("GOOGLE_API_KEY", "")
+  monkeypatch.setenv("GEMINI_API_KEY", "")
 
   # Explicit api_key takes precedence over implicit project/location.
   client = Client(vertexai=True, api_key=api_key)

google/genai/tests/client/test_client_requests.py

@@ -23,7 +23,7 @@
 
 
 def build_test_client(monkeypatch):
-  monkeypatch.setenv('GOOGLE_API_KEY', 'google_api_key')
+  monkeypatch.setenv('GEMINI_API_KEY', 'gemini_api_key')
   return Client()
 
 

google/genai/tests/conftest.py

@@ -91,7 +91,7 @@ def client(use_vertex, replays_prefix,http_options, request):
     mode = 'replay'
 
     # Set various environment variables to ensure that the test runs.
-    os.environ['GOOGLE_API_KEY'] = 'dummy-api-key'
+    os.environ['GEMINI_API_KEY'] = 'dummy-api-key'
     os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = os.path.join(
         os.path.dirname(__file__),
         'credentials.json',

google/genai/tests/live/test_live.py

@@ -81,7 +81,7 @@ async def _async_iterator_to_list(async_iter):
 
 def test_mldev_from_env(monkeypatch):
   api_key = 'google_api_key'
-  monkeypatch.setenv('GOOGLE_API_KEY', api_key)
+  monkeypatch.setenv('GEMINI_API_KEY', api_key)
 
   client = Client()
 

google/genai/tests/types/test_types.py

@@ -2215,8 +2215,8 @@ def func_under_test15() -> MyClass:
 
 
 def test_function_with_options_gemini_api(monkeypatch):
-  api_key = 'google_api_key'
-  monkeypatch.setenv('GOOGLE_API_KEY', api_key)
+  api_key = 'gemini_api_key'
+  monkeypatch.setenv('GEMINI_API_KEY', api_key)
 
   def func_under_test(a: int) -> str:
     """test return type."""