22
33import os
44import contextlib
5+ import inspect
56import dataclasses
67import pathlib
7- import types
88from typing import Any , cast
99from collections .abc import Sequence
1010import httplib2
3030 __version__ = "0.0.0"
3131
3232USER_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+ ##################
3348GENAI_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