Skip to content

Commit

Permalink
Fix default_metadata for files API. (#548)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
MarkDaoust authored Sep 11, 2024
1 parent 1ccbc5a commit 90f8094
Showing 1 changed file with 29 additions and 12 deletions.
41 changes: 29 additions & 12 deletions google/generativeai/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@

import os
import contextlib
import inspect
import dataclasses
import pathlib
import types
from typing import Any, cast
from collections.abc import Sequence
import httplib2
Expand All @@ -30,6 +30,21 @@
__version__ = "0.0.0"

USER_AGENT = "genai-py"

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


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

def _setup_discovery_api(self):
def _setup_discovery_api(self, metadata: dict | Sequence[tuple[str, str]] = ()):
api_key = self._client_options.api_key
if api_key is None:
raise ValueError(
Expand All @@ -61,6 +76,7 @@ def _setup_discovery_api(self):
http=httplib2.Http(),
postproc=lambda resp, content: (resp, content),
uri=f"{GENAI_API_DISCOVERY_URL}?version=v1beta&key={api_key}",
headers=dict(metadata),
)
response, content = request.execute()
request.http.close()
Expand All @@ -78,9 +94,10 @@ def create_file(
name: str | None = None,
display_name: str | None = None,
resumable: bool = True,
metadata: Sequence[tuple[str, str]] = (),
) -> protos.File:
if self._discovery_api is None:
self._setup_discovery_api()
self._setup_discovery_api(metadata)

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

return self.get_file({"name": result["file"]["name"]})
Expand Down Expand Up @@ -226,16 +245,14 @@ def make_client(self, name):
def keep(name, f):
if name.startswith("_"):
return False
elif name == "create_file":
return False
elif not isinstance(f, types.FunctionType):
return False
elif isinstance(f, classmethod):

if not callable(f):
return False
elif isinstance(f, staticmethod):

if "metadata" not in inspect.signature(f).parameters.keys():
return False
else:
return True

return True

def add_default_metadata_wrapper(f):
def call(*args, metadata=(), **kwargs):
Expand All @@ -244,7 +261,7 @@ def call(*args, metadata=(), **kwargs):

return call

for name, value in cls.__dict__.items():
for name, value in inspect.getmembers(cls):
if not keep(name, value):
continue
f = getattr(client, name)
Expand Down

0 comments on commit 90f8094

Please sign in to comment.