Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: introduce GEMINI_API_KEY env variable for API key, will take precedence over GOOGLE_API_KEY. #521

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down
5 changes: 4 additions & 1 deletion google/genai/_api_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 3 additions & 2 deletions google/genai/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
60 changes: 55 additions & 5 deletions google/genai/tests/client/test_client_initialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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

Expand All @@ -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()


Expand Down Expand Up @@ -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)
Expand All @@ -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
)
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion google/genai/tests/client/test_client_requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()


Expand Down
2 changes: 1 addition & 1 deletion google/genai/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
2 changes: 1 addition & 1 deletion google/genai/tests/live/test_live.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
4 changes: 2 additions & 2 deletions google/genai/tests/types/test_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down