Skip to content

Commit c9d624f

Browse files
sararobcopybara-github
sararob
authored and
copybara-github
committed
chore: add GenAI client
PiperOrigin-RevId: 755949545
1 parent acc301a commit c9d624f

7 files changed

Lines changed: 2336 additions & 6 deletions

File tree

‎setup.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@
8484
prediction_extra_require = [
8585
"docker >= 5.0.3",
8686
"fastapi >= 0.71.0, <=0.114.0",
87-
"httpx >=0.23.0, <0.25.0", # Optional dependency of fastapi
87+
"httpx >=0.23.0, <=0.28.1", # Optional dependency of fastapi
8888
"starlette >= 0.17.1",
8989
"uvicorn[standard] >= 0.16.0",
9090
]
@@ -103,7 +103,7 @@
103103
ray_extra_require = [
104104
# Cluster only supports 2.9.3, 2.33.0, and 2.42.0. Keep 2.4.0 for our
105105
# testing environment.
106-
# Note that testing is submiting a job in a cluster with Ray 2.9.3 remotely.
106+
# Note that testing is submitting a job in a cluster with Ray 2.9.3 remotely.
107107
(
108108
"ray[default] >= 2.4, <= 2.42.0,!= 2.5.*,!= 2.6.*,!= 2.7.*,!="
109109
" 2.8.*,!=2.9.0,!=2.9.1,!=2.9.2, !=2.10.*, !=2.11.*, !=2.12.*, !=2.13.*, !="
@@ -258,8 +258,9 @@
258258
"scikit-learn<1.6.0; python_version<='3.10'",
259259
"scikit-learn; python_version>'3.10'",
260260
# Lazy import requires > 2.12.0
261-
"tensorflow == 2.13.0; python_version<='3.11'",
262-
"tensorflow == 2.16.1; python_version>'3.11'",
261+
"tensorflow == 2.14.1; python_version<='3.11'",
262+
"tensorflow == 2.19.0; python_version>'3.11'",
263+
"protobuf <= 5.29.4",
263264
# TODO(jayceeli) torch 2.1.0 has conflict with pyfakefs, will check if
264265
# future versions fix this issue
265266
"torch >= 2.0.0, < 2.1.0; python_version<='3.11'",
@@ -304,6 +305,7 @@
304305
"google-cloud-bigquery >= 1.15.0, < 4.0.0, !=3.20.0",
305306
"google-cloud-resource-manager >= 1.3.3, < 3.0.0",
306307
"shapely < 3.0.0",
308+
"google-genai >= 1.0.0, <2.0.0",
307309
)
308310
+ genai_requires,
309311
extras_require={

‎testing/constraints-3.12.txt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,11 @@
44
google-api-core==2.21.0 # Tests google-api-core with rest async support
55
google-auth==2.35.0 # Tests google-auth with rest async support
66
proto-plus
7-
protobuf>=6
87
mock==4.0.2
98
google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
109
packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
1110
pytest-xdist==3.3.1 # Pinned to unbreak unit tests
1211
ray==2.5.0 # Pinned until 2.9.3 is verified for Ray tests
1312
ipython==8.22.2 # Pinned to unbreak TypeAliasType import error
14-
google-adk==0.0.2
13+
google-adk==0.0.2
14+
google-genai>=1.10.0

‎tests/unit/vertexai/genai/test_genai_client.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
# -*- coding: utf-8 -*-
2+
3+
# Copyright 2025 Google LLC
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
# pylint: disable=protected-access,bad-continuation
19+
import pytest
20+
import importlib
21+
from unittest import mock
22+
from google.cloud import aiplatform
23+
import vertexai
24+
from google.cloud.aiplatform import initializer as aiplatform_initializer
25+
26+
from vertexai import _genai
27+
28+
_TEST_PROJECT = "test-project"
29+
_TEST_LOCATION = "us-central1"
30+
31+
32+
pytestmark = pytest.mark.usefixtures("google_auth_mock")
33+
34+
35+
class TestGenAiClient:
36+
"""Unit tests for the GenAI client."""
37+
38+
def setup_method(self):
39+
importlib.reload(aiplatform_initializer)
40+
importlib.reload(aiplatform)
41+
importlib.reload(vertexai)
42+
vertexai.init(
43+
project=_TEST_PROJECT,
44+
location=_TEST_LOCATION,
45+
)
46+
47+
@pytest.mark.usefixtures("google_auth_mock")
48+
def test_genai_client(self):
49+
test_client = _genai.client.Client(
50+
project=_TEST_PROJECT, location=_TEST_LOCATION
51+
)
52+
assert test_client is not None
53+
assert test_client._api_client.vertexai
54+
assert test_client._api_client.project == _TEST_PROJECT
55+
assert test_client._api_client.location == _TEST_LOCATION
56+
57+
@pytest.mark.usefixtures("google_auth_mock")
58+
def test_evaluate_instances(self):
59+
test_client = _genai.client.Client(
60+
project=_TEST_PROJECT, location=_TEST_LOCATION
61+
)
62+
with mock.patch.object(
63+
test_client.evals, "_evaluate_instances"
64+
) as mock_evaluate:
65+
test_client.evals._evaluate_instances(bleu_input=_genai.types.BleuInput())
66+
mock_evaluate.assert_called_once_with(bleu_input=_genai.types.BleuInput())
67+
68+
@pytest.mark.usefixtures("google_auth_mock")
69+
def test_eval_run(self):
70+
test_client = _genai.client.Client(
71+
project=_TEST_PROJECT, location=_TEST_LOCATION
72+
)
73+
with pytest.raises(NotImplementedError):
74+
test_client.evals.run()

‎vertexai/_genai/__init__.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
"""The vertexai module."""
16+
17+
from . import evals
18+
from .client import Client
19+
20+
__all__ = [
21+
"Client",
22+
"evals",
23+
]

‎vertexai/_genai/client.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
16+
from typing import Optional, Union
17+
18+
import google.auth
19+
from google.genai import client
20+
from google.genai import types
21+
from .evals import Evals
22+
from .evals import AsyncEvals
23+
24+
25+
class AsyncClient:
26+
"""Async Client for the GenAI SDK."""
27+
28+
def __init__(self, api_client: client.Client):
29+
self._api_client = api_client
30+
self._aio = AsyncClient(self._api_client)
31+
self._evals = AsyncEvals(self._api_client)
32+
33+
@property
34+
def evals(self) -> AsyncEvals:
35+
return self._evals
36+
37+
38+
class Client:
39+
"""Client for the GenAI SDK.
40+
41+
Use this client to interact with Vertex-specific Gemini features.
42+
"""
43+
44+
def __init__(
45+
self,
46+
*,
47+
credentials: Optional[google.auth.credentials.Credentials] = None,
48+
project: Optional[str] = None,
49+
location: Optional[str] = None,
50+
debug_config: Optional[client.DebugConfig] = None,
51+
http_options: Optional[Union[types.HttpOptions, types.HttpOptionsDict]] = None,
52+
):
53+
"""Initializes the client.
54+
55+
Args:
56+
credentials (google.auth.credentials.Credentials): The credentials to use
57+
for authentication when calling the Vertex AI APIs. Credentials can be
58+
obtained from environment variables and default credentials. For more
59+
information, see `Set up Application Default Credentials
60+
<https://cloud.google.com/docs/authentication/provide-credentials-adc>`_.
61+
project (str): The `Google Cloud project ID
62+
<https://cloud.google.com/vertex-ai/docs/start/cloud-environment>`_ to
63+
use for quota. Can be obtained from environment variables (for example,
64+
``GOOGLE_CLOUD_PROJECT``).
65+
location (str): The `location
66+
<https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations>`_
67+
to send API requests to (for example, ``us-central1``). Can be obtained
68+
from environment variables.
69+
debug_config (DebugConfig): Config settings that control network behavior
70+
of the client. This is typically used when running test code.
71+
http_options (Union[HttpOptions, HttpOptionsDict]): Http options to use
72+
for the client.
73+
"""
74+
75+
self._debug_config = debug_config or client.DebugConfig()
76+
if isinstance(http_options, dict):
77+
http_options = types.HttpOptions(**http_options)
78+
79+
self._api_client = client.Client._get_api_client(
80+
vertexai=True,
81+
credentials=credentials,
82+
project=project,
83+
location=location,
84+
debug_config=self._debug_config,
85+
http_options=http_options,
86+
)
87+
88+
self._evals = Evals(self._api_client)
89+
90+
@property
91+
def evals(self) -> Evals:
92+
return self._evals

0 commit comments

Comments
 (0)