Function post_import_logic

patches/sitecustomize.py:77–116

(module)

@wrapt.when_imported('google.generativeai')
def post_import_logic(module):
    if os.getenv('KAGGLE_DISABLE_GOOGLE_GENERATIVE_AI_INTEGRATION') != None:
      return
    if (os.getenv('KAGGLE_DATA_PROXY_TOKEN') == None or 
       os.getenv('KAGGLE_USER_SECRETS_TOKEN') == None or
       (os.getenv('KAGGLE_DATA_PROXY_URL') == None and
       os.getenv('KAGGLE_GRPC_DATA_PROXY_URL') == None)):
      return

    old_configure = module.configure

    def new_configure(*args, **kwargs):
        if ('default_metadata' in kwargs):
            default_metadata = kwargs['default_metadata']
        else:
            default_metadata = []
        default_metadata.append(("x-kaggle-proxy-data", os.environ['KAGGLE_DATA_PROXY_TOKEN']))
        user_secrets_token = os.environ['KAGGLE_USER_SECRETS_TOKEN']
        default_metadata.append(('x-kaggle-authorization', f'Bearer {user_secrets_token}'))
        kwargs['default_metadata'] = default_metadata

        if ('client_options' in kwargs):
            client_options = kwargs['client_options']
        else:
            client_options = {}

        if os.getenv('KAGGLE_GOOGLE_GENERATIVE_AI_USE_REST_ONLY') != None:
            kwargs['transport'] = 'rest'
            
        if 'transport' in kwargs and kwargs['transport'] == 'rest':
            client_options['api_endpoint'] = os.environ['KAGGLE_DATA_PROXY_URL']
            client_options['api_endpoint'] += '/palmapi'
        else:
            client_options['api_endpoint'] = os.environ['KAGGLE_GRPC_DATA_PROXY_URL']
        kwargs['client_options'] = client_options

        old_configure(*args, **kwargs)

    module.configure = new_configure
    module.configure() # generativeai can use GOOGLE_API_KEY env variable, so make sure we have the other configs set

@wrapt.when_imported('google.genai')
def post_genai_import_logic(module):