Skip to content

Commit 3a38b50

Browse files
wanlin31copybara-github
authored andcommitted
feat: introduce GEMINI_API_KEY env variable for API key, will take precedence over GOOGLE_API_KEY.
PiperOrigin-RevId: 737771755
1 parent f335107 commit 3a38b50

8 files changed

+71
-15
lines changed

README.md

+4-2
Original file line numberDiff line numberDiff line change
@@ -45,10 +45,12 @@ You can create a client by configuring the necessary environment variables.
4545
Configuration setup instructions depends on whether you're using the Gemini API
4646
on Vertex AI or the ML Dev Gemini API.
4747

48-
**ML Dev Gemini API:** Set `GOOGLE_API_KEY` as shown below:
48+
**ML Dev Gemini API:** Set `GEMINI_API_KEY`(preferred) or `GOOGLE_API_KEY` as
49+
shown below:
4950

5051
```bash
51-
export GOOGLE_API_KEY='your-api-key'
52+
export GEMINI_API_KEY='your-api-key'
53+
5254
```
5355

5456
**Vertex AI API:** Set `GOOGLE_GENAI_USE_VERTEXAI`, `GOOGLE_CLOUD_PROJECT`

google/genai/_api_client.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,10 @@ def __init__(
287287
# Retrieve implicitly set values from the environment.
288288
env_project = os.environ.get('GOOGLE_CLOUD_PROJECT', None)
289289
env_location = os.environ.get('GOOGLE_CLOUD_LOCATION', None)
290-
env_api_key = os.environ.get('GOOGLE_API_KEY', None)
290+
# Gemini API key takes precedence over Google API key.
291+
env_api_key = os.environ.get('GEMINI_API_KEY', None)
292+
if not env_api_key or env_api_key == '':
293+
env_api_key = os.environ.get('GOOGLE_API_KEY', None)
291294
self.project = project or env_project
292295
self.location = location or env_location
293296
self.api_key = api_key or env_api_key

google/genai/client.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -102,8 +102,9 @@ class Client:
102102
To initialize the client, provide the required arguments either directly
103103
or by using environment variables. Gemini API users and Vertex AI users in
104104
express mode can provide API key by providing input argument
105-
`api_key="your-api-key"` or by defining `GOOGLE_API_KEY="your-api-key"` as an
106-
environment variable
105+
`api_key="your-api-key"` or via environment variables: `GOOGLE_API_KEY` or
106+
`GEMINI_API_KEY`. When both environment variables are set `GEMINI_API_KEY`
107+
takes precedence.
107108
108109
Vertex AI API users can provide inputs argument as `vertexai=True,
109110
project="your-project-id", location="us-central1"` or by defining

google/genai/tests/client/test_client_initialization.py

+55-5
Original file line numberDiff line numberDiff line change
@@ -26,16 +26,41 @@
2626
from ... import Client
2727

2828

29-
def test_ml_dev_from_env(monkeypatch):
29+
def test_ml_dev_from_google_api_key_env(monkeypatch):
3030
api_key = "google_api_key"
3131
monkeypatch.setenv("GOOGLE_API_KEY", api_key)
32+
monkeypatch.setenv("GEMINI_API_KEY", "")
3233

3334
client = Client()
3435

3536
assert not client.models._api_client.vertexai
3637
assert client.models._api_client.api_key == api_key
3738
assert isinstance(client.models._api_client, api_client.BaseApiClient)
3839

40+
def test_ml_dev_from_gemini_api_key_env(monkeypatch):
41+
api_key = "gemini_api_key"
42+
monkeypatch.setenv("GOOGLE_API_KEY", "")
43+
monkeypatch.setenv("GEMINI_API_KEY", api_key)
44+
45+
client = Client()
46+
47+
assert not client.models._api_client.vertexai
48+
assert client.models._api_client.api_key == api_key
49+
assert isinstance(client.models._api_client, api_client.BaseApiClient)
50+
51+
def test_ml_dev_from_env_precedence(monkeypatch):
52+
# Gemini API key takes precedence over Google API key.
53+
gemini_api_key = "gemini_api_key"
54+
monkeypatch.setenv("GEMINI_API_KEY", gemini_api_key)
55+
google_api_key = "google_api_key"
56+
monkeypatch.setenv("GOOGLE_API_KEY", google_api_key)
57+
58+
client = Client()
59+
60+
assert not client.models._api_client.vertexai
61+
assert client.models._api_client.api_key == gemini_api_key
62+
assert isinstance(client.models._api_client, api_client.BaseApiClient)
63+
3964

4065
def test_ml_dev_from_constructor():
4166
api_key = "google_api_key"
@@ -271,6 +296,7 @@ def test_invalid_vertexai_constructor_empty(monkeypatch):
271296
monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "")
272297
monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "")
273298
monkeypatch.setenv("GOOGLE_API_KEY", "")
299+
monkeypatch.setenv("GEMINI_API_KEY", "")
274300
def mock_auth_default(scopes=None):
275301
return None, None
276302

@@ -281,6 +307,7 @@ def mock_auth_default(scopes=None):
281307
def test_invalid_mldev_constructor_empty(monkeypatch):
282308
with pytest.raises(ValueError):
283309
monkeypatch.setenv("GOOGLE_API_KEY", "")
310+
monkeypatch.setenv("GEMINI_API_KEY", "")
284311
Client()
285312

286313

@@ -389,8 +416,8 @@ def test_mldev_explicit_arg_precedence(monkeypatch):
389416

390417

391418
def test_replay_client_ml_dev_from_env(monkeypatch, use_vertex: bool):
392-
api_key = "google_api_key"
393-
monkeypatch.setenv("GOOGLE_API_KEY", api_key)
419+
gemini_api_key = "gemini_api_key"
420+
monkeypatch.setenv("GEMINI_API_KEY", gemini_api_key)
394421
monkeypatch.setenv("GOOGLE_GENAI_CLIENT_MODE", "replay")
395422
api_type = "vertex" if use_vertex else "mldev"
396423
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):
399426
client = Client()
400427

401428
assert not client.models._api_client.vertexai
402-
assert client.models._api_client.api_key == api_key
429+
assert client.models._api_client.api_key == gemini_api_key
403430
assert isinstance(
404431
client.models._api_client, replay_api_client.ReplayApiClient
405432
)
@@ -461,10 +488,11 @@ def test_vertexai_apikey_from_constructor(monkeypatch):
461488
assert isinstance(client.models._api_client, api_client.BaseApiClient)
462489

463490

464-
def test_vertexai_apikey_from_env(monkeypatch):
491+
def test_vertexai_apikey_from_google_api_key_env(monkeypatch):
465492
# Vertex AI Express mode uses API key on Vertex AI.
466493
api_key = "vertexai_api_key"
467494
monkeypatch.setenv("GOOGLE_API_KEY", api_key)
495+
monkeypatch.setenv("GEMINI_API_KEY", "")
468496

469497
# Due to proj/location taking precedence, need to clear proj/location env
470498
# variables.
@@ -480,6 +508,27 @@ def test_vertexai_apikey_from_env(monkeypatch):
480508
assert "aiplatform" in client._api_client._http_options["base_url"]
481509
assert isinstance(client.models._api_client, api_client.BaseApiClient)
482510

511+
def test_vertexai_apikey_from_env_precedence(monkeypatch):
512+
# Vertex AI Express mode uses gemnai API key on Vertex AI.
513+
gemini_api_key = "gemini_api_key"
514+
monkeypatch.setenv("GEMINI_API_KEY", gemini_api_key)
515+
google_api_key = "google_api_key"
516+
monkeypatch.setenv("GOOGLE_API_KEY", google_api_key)
517+
518+
# Due to proj/location taking precedence, need to clear proj/location env
519+
# variables.
520+
monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "")
521+
monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "")
522+
523+
client = Client(vertexai=True)
524+
525+
assert client.models._api_client.vertexai
526+
assert client.models._api_client.api_key == gemini_api_key
527+
assert not client.models._api_client.project
528+
assert not client.models._api_client.location
529+
assert "aiplatform" in client._api_client._http_options["base_url"]
530+
assert isinstance(client.models._api_client, api_client.BaseApiClient)
531+
483532

484533
def test_vertexai_apikey_invalid_constructor1():
485534
# Vertex AI Express mode uses API key on Vertex AI.
@@ -504,6 +553,7 @@ def test_vertexai_apikey_combo1(monkeypatch):
504553
monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", project_id)
505554
monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", location)
506555
monkeypatch.setenv("GOOGLE_API_KEY", "")
556+
monkeypatch.setenv("GEMINI_API_KEY", "")
507557

508558
# Explicit api_key takes precedence over implicit project/location.
509559
client = Client(vertexai=True, api_key=api_key)

google/genai/tests/client/test_client_requests.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323

2424

2525
def build_test_client(monkeypatch):
26-
monkeypatch.setenv('GOOGLE_API_KEY', 'google_api_key')
26+
monkeypatch.setenv('GEMINI_API_KEY', 'gemini_api_key')
2727
return Client()
2828

2929

google/genai/tests/conftest.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ def client(use_vertex, replays_prefix,http_options, request):
9191
mode = 'replay'
9292

9393
# Set various environment variables to ensure that the test runs.
94-
os.environ['GOOGLE_API_KEY'] = 'dummy-api-key'
94+
os.environ['GEMINI_API_KEY'] = 'dummy-api-key'
9595
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = os.path.join(
9696
os.path.dirname(__file__),
9797
'credentials.json',

google/genai/tests/live/test_live.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ async def _async_iterator_to_list(async_iter):
8181

8282
def test_mldev_from_env(monkeypatch):
8383
api_key = 'google_api_key'
84-
monkeypatch.setenv('GOOGLE_API_KEY', api_key)
84+
monkeypatch.setenv('GEMINI_API_KEY', api_key)
8585

8686
client = Client()
8787

google/genai/tests/types/test_types.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -2215,8 +2215,8 @@ def func_under_test15() -> MyClass:
22152215

22162216

22172217
def test_function_with_options_gemini_api(monkeypatch):
2218-
api_key = 'google_api_key'
2219-
monkeypatch.setenv('GOOGLE_API_KEY', api_key)
2218+
api_key = 'gemini_api_key'
2219+
monkeypatch.setenv('GEMINI_API_KEY', api_key)
22202220

22212221
def func_under_test(a: int) -> str:
22222222
"""test return type."""

0 commit comments

Comments
 (0)