use the normal dataproxy if REST only

Kaggle · Dec 5, 2023 · e23d819 · e23d819
1 parent 2c1d3a0
commit e23d819
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/patches/sitecustomize.py b/patches/sitecustomize.py
@@ -105,6 +105,7 @@ def new_configure(*args, **kwargs):
         if os.getenv('KAGGLE_GOOGLE_GENERATIVE_AI_USE_REST_ONLY') != None:
             client_options['api_endpoint'] += '/palmapi'
             kwargs['transport'] = 'rest'
+            client_options['api_endpoint'] = os.environ['KAGGLE_DATA_PROXY_URL']
         elif 'transport' in kwargs and kwargs['transport'] == 'rest':
             client_options['api_endpoint'] += '/palmapi'
         kwargs['client_options'] = client_options

diff --git a/tests/test_google_generativeai_patch.py b/tests/test_google_generativeai_patch.py
@@ -32,7 +32,8 @@ def test_proxy_enabled(self):
         proxy_token = "proxy_token"
         env.set("KAGGLE_USER_SECRETS_TOKEN", secrets_token)
         env.set("KAGGLE_DATA_PROXY_TOKEN", proxy_token)
-        env.set("KAGGLE_GRPC_DATA_PROXY_URL", self.endpoint)
+        env.set("KAGGLE_DATA_PROXY_URL", self.endpoint)
+        env.set("KAGGLE_GRPC_DATA_PROXY_URL", "http://127.0.0.1:50001")
         env.set("KAGGLE_GOOGLE_GENERATIVE_AI_USE_REST_ONLY", "True")
         server_address = urlparse(self.endpoint)
         with env: