Skip to content

Commit 90f8094

Browse files
authored
Fix default_metadata for files API. (#548)
* Fix default_metadata for files API. Change-Id: Ibba74ccda137fe4ac62525868e036a0a0d612348 * Pass the metadata through to the discovert URL Change-Id: If8f938a1026d46aafc617d1c0e95360c463f0ba0 * fix tests, getmembers doesn't return classmethod/staticmetnod objects. Change-Id: I9449c6f3299f820fea4128cb598f39cf359cd9ef * format Change-Id: I4133647fbea3b6b995ed6edf56bd2a243c2131af * Add notes describing interaction of the discovery URL and API endpoint. Change-Id: Ibb8740fd548cc8a8a75b32508d6d05aa6e89f633 * note abput default-metadata Change-Id: I9647c47718d7b6981692b484b9e9b2ed0f9a42cc * remove unused import Change-Id: Ie49216fef5eb13e202fd276fff62e7e20f03b062
1 parent 1ccbc5a commit 90f8094

File tree

1 file changed

+29
-12
lines changed

1 file changed

+29
-12
lines changed

google/generativeai/client.py

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22

33
import os
44
import contextlib
5+
import inspect
56
import dataclasses
67
import pathlib
7-
import types
88
from typing import Any, cast
99
from collections.abc import Sequence
1010
import httplib2
@@ -30,6 +30,21 @@
3030
__version__ = "0.0.0"
3131

3232
USER_AGENT = "genai-py"
33+
34+
#### Caution! ####
35+
# - It would make sense for the discovery URL to respect the client_options.endpoint setting.
36+
# - That would make testing Files on the staging server possible.
37+
# - We tried fixing this once, but broke colab in the process because their endpoint didn't forward the discovery
38+
# requests. https://github.com/google-gemini/generative-ai-python/pull/333
39+
# - Kaggle would have a similar problem (b/362278209).
40+
# - I think their proxy would forward the discovery traffic.
41+
# - But they don't need to intercept the files-service at all, and uploads of large files could overload them.
42+
# - Do the scotty uploads go to the same domain?
43+
# - If you do route the discovery call to kaggle, be sure to attach the default_metadata (they need it).
44+
# - One solution to all this would be if configure could take overrides per service.
45+
# - set client_options.endpoint, but use a different endpoint for file service? It's not clear how best to do that
46+
# through the file service.
47+
##################
3348
GENAI_API_DISCOVERY_URL = "https://generativelanguage.googleapis.com/$discovery/rest"
3449

3550

@@ -50,7 +65,7 @@ def __init__(self, *args, **kwargs):
5065
self._discovery_api = None
5166
super().__init__(*args, **kwargs)
5267

53-
def _setup_discovery_api(self):
68+
def _setup_discovery_api(self, metadata: dict | Sequence[tuple[str, str]] = ()):
5469
api_key = self._client_options.api_key
5570
if api_key is None:
5671
raise ValueError(
@@ -61,6 +76,7 @@ def _setup_discovery_api(self):
6176
http=httplib2.Http(),
6277
postproc=lambda resp, content: (resp, content),
6378
uri=f"{GENAI_API_DISCOVERY_URL}?version=v1beta&key={api_key}",
79+
headers=dict(metadata),
6480
)
6581
response, content = request.execute()
6682
request.http.close()
@@ -78,9 +94,10 @@ def create_file(
7894
name: str | None = None,
7995
display_name: str | None = None,
8096
resumable: bool = True,
97+
metadata: Sequence[tuple[str, str]] = (),
8198
) -> protos.File:
8299
if self._discovery_api is None:
83-
self._setup_discovery_api()
100+
self._setup_discovery_api(metadata)
84101

85102
file = {}
86103
if name is not None:
@@ -92,6 +109,8 @@ def create_file(
92109
filename=path, mimetype=mime_type, resumable=resumable
93110
)
94111
request = self._discovery_api.media().upload(body={"file": file}, media_body=media)
112+
for key, value in metadata:
113+
request.headers[key] = value
95114
result = request.execute()
96115

97116
return self.get_file({"name": result["file"]["name"]})
@@ -226,16 +245,14 @@ def make_client(self, name):
226245
def keep(name, f):
227246
if name.startswith("_"):
228247
return False
229-
elif name == "create_file":
230-
return False
231-
elif not isinstance(f, types.FunctionType):
232-
return False
233-
elif isinstance(f, classmethod):
248+
249+
if not callable(f):
234250
return False
235-
elif isinstance(f, staticmethod):
251+
252+
if "metadata" not in inspect.signature(f).parameters.keys():
236253
return False
237-
else:
238-
return True
254+
255+
return True
239256

240257
def add_default_metadata_wrapper(f):
241258
def call(*args, metadata=(), **kwargs):
@@ -244,7 +261,7 @@ def call(*args, metadata=(), **kwargs):
244261

245262
return call
246263

247-
for name, value in cls.__dict__.items():
264+
for name, value in inspect.getmembers(cls):
248265
if not keep(name, value):
249266
continue
250267
f = getattr(client, name)

0 commit comments

Comments
 (0)