feat(python): Support GEMINI_API_KEY as environment variable for setting API key.

PiperOrigin-RevId: 767169011

Author: wanlin31

google/genai/_api_client.py

@@ -68,6 +68,26 @@ class EphemeralTokenAPIKeyError(ValueError):
   """Error raised when the API key is invalid."""
 
 
+# This method checks for the API key in the environment variables. Google API
+# key is precedenced over Gemini API key.
+def _get_env_api_key() -> Optional[str]:
+  """Gets the API key from environment variables, prioritizing GOOGLE_API_KEY.
+
+  Returns:
+      The API key string if found, otherwise None. Empty string is considered
+      invalid.
+  """
+  env_google_api_key = os.environ.get('GOOGLE_API_KEY', None)
+  env_gemini_api_key = os.environ.get('GEMINI_API_KEY', None)
+  if env_google_api_key and env_gemini_api_key:
+    logger.warning(
+        'Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using'
+        ' GOOGLE_API_KEY.'
+    )
+
+  return env_google_api_key or env_gemini_api_key or None
+
+
 def _append_library_version_headers(headers: dict[str, str]) -> None:
   """Appends the telemetry header to the headers dict."""
   library_label = f'google-genai-sdk/{version.__version__}'
@@ -371,7 +391,7 @@ 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)
+    env_api_key = _get_env_api_key()
     self.project = project or env_project
     self.location = location or env_location
     self.api_key = api_key or env_api_key

google/genai/tests/client/test_client_initialization.py

@@ -31,9 +31,34 @@
 from ... import Client
 
 
-def test_ml_dev_from_env(monkeypatch):
+def test_ml_dev_from_gemini_env_only(monkeypatch):
+  api_key = "gemini_api_key"
+  monkeypatch.setenv("GEMINI_API_KEY", api_key)
+  monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
+
+  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_env_with_google_env_empty(monkeypatch):
+  api_key = "gemini_api_key"
+  monkeypatch.setenv("GEMINI_API_KEY", api_key)
+  monkeypatch.setenv("GOOGLE_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_google_env_only(monkeypatch):
   api_key = "google_api_key"
   monkeypatch.setenv("GOOGLE_API_KEY", api_key)
+  monkeypatch.delenv("GEMINI_API_KEY", raising=False)
 
   client = Client()
 
@@ -42,6 +67,24 @@ def test_ml_dev_from_env(monkeypatch):
   assert isinstance(client.models._api_client, api_client.BaseApiClient)
 
 
+def test_ml_dev_both_env_key_set(monkeypatch, caplog):
+  caplog.set_level(logging.DEBUG, logger="google_genai._api_client")
+  google_api_key = "google_api_key"
+  gemini_api_key = "gemini_api_key"
+  monkeypatch.setenv("GOOGLE_API_KEY", google_api_key)
+  monkeypatch.setenv("GEMINI_API_KEY", gemini_api_key)
+
+  client = Client()
+
+  assert not client.models._api_client.vertexai
+  assert client.models._api_client.api_key == google_api_key
+  assert isinstance(client.models._api_client, api_client.BaseApiClient)
+  assert (
+      "Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using GOOGLE_API_KEY."
+      in caplog.text
+  )
+
+
 def test_ml_dev_from_constructor():
   api_key = "google_api_key"
 
@@ -276,6 +319,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
@@ -287,6 +331,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()
 
 
@@ -379,16 +424,22 @@ def test_invalid_mldev_constructor():
     assert isinstance(e, ValueError)
 
 
-def test_mldev_explicit_arg_precedence(monkeypatch):
+def test_mldev_explicit_arg_precedence(monkeypatch, caplog):
+  caplog.set_level(logging.DEBUG, logger="google_genai._api_client")
   api_key = "constructor_api_key"
 
-  monkeypatch.setenv("GOOGLE_API_KEY", "env_api_key")
+  monkeypatch.setenv("GOOGLE_API_KEY", "google_env_api_key")
+  monkeypatch.setenv("GEMINI_API_KEY", "gemini_env_api_key")
 
   client = Client(api_key=api_key)
 
   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)
+  assert (
+      "Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using GOOGLE_API_KEY."
+      in caplog.text
+  )
 
 
 def test_replay_client_ml_dev_from_env(monkeypatch, use_vertex: bool):
@@ -464,10 +515,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_env_google_api_key_only(monkeypatch):
   # Vertex AI Express mode uses API key on Vertex AI.
   api_key = "vertexai_api_key"
   monkeypatch.setenv("GOOGLE_API_KEY", api_key)
+  monkeypatch.delenv("GEMINI_API_KEY", raising=False)
 
   # Due to proj/location taking precedence, need to clear proj/location env
   # variables.
@@ -484,6 +536,75 @@ def test_vertexai_apikey_from_env(monkeypatch):
   assert isinstance(client.models._api_client, api_client.BaseApiClient)
 
 
+def test_vertexai_apikey_from_env_gemini_api_key_only(monkeypatch):
+  # Vertex AI Express mode uses API key on Vertex AI.
+  api_key = "vertexai_api_key"
+  monkeypatch.setenv("GEMINI_API_KEY", api_key)
+  monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
+
+  # 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 == 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_from_env_gemini_api_key_with_google_api_key_empty(monkeypatch):
+  # Vertex AI Express mode uses API key on Vertex AI.
+  api_key = "vertexai_api_key"
+  monkeypatch.setenv("GEMINI_API_KEY", api_key)
+  monkeypatch.setenv("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 == 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_from_env_both_api_keys(monkeypatch,  caplog):
+  caplog.set_level(logging.DEBUG, logger="google_genai._api_client")
+  # Vertex AI Express mode uses API key on Vertex AI.
+  google_api_key = "google_api_key"
+  gemini_api_key = "vertexai_api_key"
+  monkeypatch.setenv("GEMINI_API_KEY", gemini_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 == google_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)
+  assert (
+      "Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using GOOGLE_API_KEY."
+      in caplog.text
+  )
+
+
 def test_vertexai_apikey_invalid_constructor1():
   # Vertex AI Express mode uses API key on Vertex AI.
   api_key = "vertexai_api_key"