Patching for generativeai Python package

Dockerfile.tmpl

@@ -634,6 +634,10 @@ RUN mkdir -p ~/src && git clone https://github.com/SohierDane/BigQuery_Helper ~/
     sed -i 's/)/packages=["bq_helper"])/g' ~/src/BigQuery_Helper/setup.py && \
     pip install -e ~/src/BigQuery_Helper && \
     /tmp/clean-layer.sh
+
+RUN pip install wrapt \
+        google-generativeai && \
+    /tmp/clean-layer.sh
 
 # Add BigQuery client proxy settings
 ENV PYTHONUSERBASE "/root/.local"

patches/sitecustomize.py

@@ -7,6 +7,8 @@
 import importlib
 import importlib.machinery
 
+import wrapt
+
 class GcpModuleFinder(importlib.abc.MetaPathFinder):
     _MODULES = [
         'google.cloud.bigquery',
@@ -73,3 +75,27 @@ def exec_module(self, module):
 
 if not hasattr(sys, 'frozen'):
     sys.meta_path.insert(0, GcpModuleFinder())
+
+@wrapt.when_imported('google.generativeai')
+def post_import_logic(module):
+    if os.getenv('KAGGLE_DISABLE_GOOGLE_GENERATIVE_AI_INTEGRATION') == 1:
+      return
+    old_configure = module.configure
+    def new_configure(*args, **kwargs):
+        if ('default_metadata' in kwargs):
+            default_metadata = kwargs['default_metadata']
+        else:
+            default_metadata = []
+        kwargs['transport'] = 'rest' # Only support REST requests for now
+        default_metadata.append(("x-kaggle-proxy-data", os.environ['KAGGLE_DATA_PROXY_TOKEN']))
+        default_metadata.append(('x-kaggle-authorization', f'Bearer {os.environ['KAGGLE_USER_SECRETS_TOKEN']}'))
+        kwargs['default_metadata'] = default_metadata
+        if ('client_options' in kwargs):
+            client_options = kwargs['client_options']
+        else:
+            client_options = {}
+        client_options['api_endpoint'] = os.environ['KAGGLE_DATA_PROXY_URL'] + '/palmapi'
+        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