|
1 | 1 | from __future__ import annotations
|
2 | 2 |
|
3 | 3 | import os
|
| 4 | +import contextlib |
4 | 5 | import dataclasses
|
5 | 6 | import pathlib
|
6 |
| -import re |
7 | 7 | import types
|
8 | 8 | from typing import Any, cast
|
9 | 9 | from collections.abc import Sequence
|
|
12 | 12 | import google.ai.generativelanguage as glm
|
13 | 13 |
|
14 | 14 | from google.auth import credentials as ga_credentials
|
| 15 | +from google.auth import exceptions as ga_exceptions |
| 16 | +from google import auth |
15 | 17 | from google.api_core import client_options as client_options_lib
|
16 | 18 | from google.api_core import gapic_v1
|
17 | 19 | from google.api_core import operations_v1
|
|
30 | 32 | GENAI_API_DISCOVERY_URL = "https://generativelanguage.googleapis.com/$discovery/rest"
|
31 | 33 |
|
32 | 34 |
|
| 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 | + |
33 | 47 | class FileServiceClient(glm.FileServiceClient):
|
34 | 48 | def __init__(self, *args, **kwargs):
|
35 | 49 | self._discovery_api = None
|
@@ -183,7 +197,17 @@ def make_client(self, name):
|
183 | 197 | if not self.client_config:
|
184 | 198 | configure()
|
185 | 199 |
|
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 |
187 | 211 |
|
188 | 212 | if not self.default_metadata:
|
189 | 213 | return client
|
|
0 commit comments