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"
@@ -647,6 +651,9 @@ ADD patches/sitecustomize.py /root/.local/lib/python3.10/site-packages/sitecusto
 # Override default imagemagick policies
 ADD patches/imagemagick-policy.xml /etc/ImageMagick-6/policy.xml
 
+# Add generativeai Python patch
+CMD python patches/generativeaipatch.py
+
 # Add Kaggle module resolver
 ADD patches/kaggle_module_resolver.py /opt/conda/lib/python3.10/site-packages/tensorflow_hub/kaggle_module_resolver.py
 RUN sed -i '/from tensorflow_hub import uncompressed_module_resolver/a from tensorflow_hub import kaggle_module_resolver' /opt/conda/lib/python3.10/site-packages/tensorflow_hub/config.py && \

patches/generativeaipatch.py

@@ -0,0 +1,24 @@
+import wrapt
+import os
+
+@wrapt.when_imported('google.generativeai')
+def post_import_logic(module):
+    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