Skip to content

Commit c6e4087

Browse files
authored
Add google-genai to docker image (#1455)
1 parent 2cb0a25 commit c6e4087

File tree

2 files changed

+56
-0
lines changed

2 files changed

+56
-0
lines changed

kaggle_requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ google-cloud-automl==1.0.1
4848
google-cloud-translate==3.12.1
4949
google-cloud-videointelligence
5050
google-cloud-vision
51+
google-genai
5152
gpxpy
5253
h2o
5354
haversine

tests/test_google_genai_patch.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import json
2+
import unittest
3+
import threading
4+
5+
from test.support.os_helper import EnvironmentVarGuard
6+
from urllib.parse import urlparse
7+
8+
from http.server import BaseHTTPRequestHandler, HTTPServer
9+
10+
class HTTPHandler(BaseHTTPRequestHandler):
11+
called = False
12+
path = None
13+
headers = {}
14+
15+
def do_HEAD(self):
16+
self.send_response(200)
17+
18+
def do_POST(self):
19+
HTTPHandler.path = self.path
20+
HTTPHandler.headers = self.headers
21+
HTTPHandler.called = True
22+
self.send_response(200)
23+
self.send_header("Content-type", "application/json")
24+
self.end_headers()
25+
26+
class TestGoogleGenAiPatch(unittest.TestCase):
27+
endpoint = "http://127.0.0.1:80"
28+
29+
def test_proxy_enabled(self):
30+
env = EnvironmentVarGuard()
31+
secrets_token = "secrets_token"
32+
proxy_token = "proxy_token"
33+
env.set("KAGGLE_USER_SECRETS_TOKEN", secrets_token)
34+
env.set("KAGGLE_DATA_PROXY_TOKEN", proxy_token)
35+
env.set("KAGGLE_DATA_PROXY_URL", self.endpoint)
36+
server_address = urlparse(self.endpoint)
37+
with env:
38+
with HTTPServer((server_address.hostname, server_address.port), HTTPHandler) as httpd:
39+
threading.Thread(target=httpd.serve_forever).start()
40+
from google import genai
41+
api_key = "NotARealAPIKey"
42+
client = genai.Client(api_key = api_key)
43+
try:
44+
client.models.generate_content(
45+
model="gemini-2.0-flash-exp",
46+
contents="What's the largest planet in our solar system?"
47+
)
48+
except:
49+
pass
50+
httpd.shutdown()
51+
self.assertTrue(HTTPHandler.called)
52+
self.assertIn("/palmapi", HTTPHandler.path)
53+
self.assertEqual(proxy_token, HTTPHandler.headers["x-kaggle-proxy-data"])
54+
self.assertEqual("Bearer {}".format(secrets_token), HTTPHandler.headers["x-kaggle-authorization"])
55+
self.assertEqual(api_key, HTTPHandler.headers["x-goog-api-key"])

0 commit comments

Comments
 (0)