Skip to content

Commit f987fde

Browse files
authored
improve the no-credentials error message, fail fast for no-credentials in colab. (#352)
* improve the no-credentials error message Change-Id: I294bd094b56287ed923716dce9ea705ef3135f5b * patch colab credentials Change-Id: I5a3cb3168448a565eb3cdc8a0063ae041c41a260 * format Change-Id: I013d506bdcb64092daddedcf3e30f3728a8f3e30
1 parent 05877f7 commit f987fde

File tree

1 file changed

+26
-2
lines changed

1 file changed

+26
-2
lines changed

google/generativeai/client.py

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
from __future__ import annotations
22

33
import os
4+
import contextlib
45
import dataclasses
56
import pathlib
6-
import re
77
import types
88
from typing import Any, cast
99
from collections.abc import Sequence
@@ -12,6 +12,8 @@
1212
import google.ai.generativelanguage as glm
1313

1414
from google.auth import credentials as ga_credentials
15+
from google.auth import exceptions as ga_exceptions
16+
from google import auth
1517
from google.api_core import client_options as client_options_lib
1618
from google.api_core import gapic_v1
1719
from google.api_core import operations_v1
@@ -30,6 +32,18 @@
3032
GENAI_API_DISCOVERY_URL = "https://generativelanguage.googleapis.com/$discovery/rest"
3133

3234

35+
@contextlib.contextmanager
36+
def patch_colab_gce_credentials():
37+
get_gce = auth._default._get_gce_credentials
38+
if "COLAB_RELEASE_TAG" in os.environ:
39+
auth._default._get_gce_credentials = lambda *args, **kwargs: (None, None)
40+
41+
try:
42+
yield
43+
finally:
44+
auth._default._get_gce_credentials = get_gce
45+
46+
3347
class FileServiceClient(glm.FileServiceClient):
3448
def __init__(self, *args, **kwargs):
3549
self._discovery_api = None
@@ -183,7 +197,17 @@ def make_client(self, name):
183197
if not self.client_config:
184198
configure()
185199

186-
client = cls(**self.client_config)
200+
try:
201+
with patch_colab_gce_credentials():
202+
client = cls(**self.client_config)
203+
except ga_exceptions.DefaultCredentialsError as e:
204+
e.args = (
205+
"\n No API_KEY or ADC found. Please either:\n"
206+
" - Set the `GOOGLE_API_KEY` environment variable.\n"
207+
" - Manually pass the key with `genai.configure(api_key=my_api_key)`.\n"
208+
" - Or set up Application Default Credentials, see https://ai.google.dev/gemini-api/docs/oauth for more information.",
209+
)
210+
raise e
187211

188212
if not self.default_metadata:
189213
return client

0 commit comments

Comments
 (0)