chore: add GenAI client

sararob · copybara-github · commit c9d624fd80b3 · 2025-05-07T12:05:26.000-07:00
PiperOrigin-RevId: 755949545
diff --git a/setup.py b/setup.py
@@ -84,7 +84,7 @@
 prediction_extra_require = [
     "docker >= 5.0.3",
     "fastapi >= 0.71.0, <=0.114.0",
-    "httpx >=0.23.0, <0.25.0",  # Optional dependency of fastapi
+    "httpx >=0.23.0, <=0.28.1",  # Optional dependency of fastapi
     "starlette >= 0.17.1",
     "uvicorn[standard] >= 0.16.0",
 ]
@@ -103,7 +103,7 @@
 ray_extra_require = [
     # Cluster only supports 2.9.3, 2.33.0, and 2.42.0. Keep 2.4.0 for our
     # testing environment.
-    # Note that testing is submiting a job in a cluster with Ray 2.9.3 remotely.
+    # Note that testing is submitting a job in a cluster with Ray 2.9.3 remotely.
     (
         "ray[default] >= 2.4, <= 2.42.0,!= 2.5.*,!= 2.6.*,!= 2.7.*,!="
         " 2.8.*,!=2.9.0,!=2.9.1,!=2.9.2, !=2.10.*, !=2.11.*, !=2.12.*, !=2.13.*, !="
@@ -258,8 +258,9 @@
         "scikit-learn<1.6.0; python_version<='3.10'",
         "scikit-learn; python_version>'3.10'",
         # Lazy import requires > 2.12.0
-        "tensorflow == 2.13.0; python_version<='3.11'",
-        "tensorflow == 2.16.1; python_version>'3.11'",
+        "tensorflow == 2.14.1; python_version<='3.11'",
+        "tensorflow == 2.19.0; python_version>'3.11'",
+        "protobuf <= 5.29.4",
         # TODO(jayceeli) torch 2.1.0 has conflict with pyfakefs, will check if
         # future versions fix this issue
         "torch >= 2.0.0, < 2.1.0; python_version<='3.11'",
@@ -304,6 +305,7 @@
         "google-cloud-bigquery >= 1.15.0, < 4.0.0, !=3.20.0",
         "google-cloud-resource-manager >= 1.3.3, < 3.0.0",
         "shapely < 3.0.0",
+        "google-genai >= 1.0.0, <2.0.0",
     )
     + genai_requires,
     extras_require={
diff --git a/testing/constraints-3.12.txt b/testing/constraints-3.12.txt
@@ -4,11 +4,11 @@
 google-api-core==2.21.0 # Tests google-api-core with rest async support
 google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus
-protobuf>=6
 mock==4.0.2
 google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
 packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.5.0 # Pinned until 2.9.3 is verified for Ray tests
 ipython==8.22.2 # Pinned to unbreak TypeAliasType import error
-google-adk==0.0.2
+google-adk==0.0.2
+google-genai>=1.10.0
diff --git a/tests/unit/vertexai/genai/test_genai_client.py b/tests/unit/vertexai/genai/test_genai_client.py
@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# pylint: disable=protected-access,bad-continuation
+import pytest
+import importlib
+from unittest import mock
+from google.cloud import aiplatform
+import vertexai
+from google.cloud.aiplatform import initializer as aiplatform_initializer
+
+from vertexai import _genai
+
+_TEST_PROJECT = "test-project"
+_TEST_LOCATION = "us-central1"
+
+
+pytestmark = pytest.mark.usefixtures("google_auth_mock")
+
+
+class TestGenAiClient:
+    """Unit tests for the GenAI client."""
+
+    def setup_method(self):
+        importlib.reload(aiplatform_initializer)
+        importlib.reload(aiplatform)
+        importlib.reload(vertexai)
+        vertexai.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+
+    @pytest.mark.usefixtures("google_auth_mock")
+    def test_genai_client(self):
+        test_client = _genai.client.Client(
+            project=_TEST_PROJECT, location=_TEST_LOCATION
+        )
+        assert test_client is not None
+        assert test_client._api_client.vertexai
+        assert test_client._api_client.project == _TEST_PROJECT
+        assert test_client._api_client.location == _TEST_LOCATION
+
+    @pytest.mark.usefixtures("google_auth_mock")
+    def test_evaluate_instances(self):
+        test_client = _genai.client.Client(
+            project=_TEST_PROJECT, location=_TEST_LOCATION
+        )
+        with mock.patch.object(
+            test_client.evals, "_evaluate_instances"
+        ) as mock_evaluate:
+            test_client.evals._evaluate_instances(bleu_input=_genai.types.BleuInput())
+            mock_evaluate.assert_called_once_with(bleu_input=_genai.types.BleuInput())
+
+    @pytest.mark.usefixtures("google_auth_mock")
+    def test_eval_run(self):
+        test_client = _genai.client.Client(
+            project=_TEST_PROJECT, location=_TEST_LOCATION
+        )
+        with pytest.raises(NotImplementedError):
+            test_client.evals.run()
diff --git a/vertexai/_genai/__init__.py b/vertexai/_genai/__init__.py
@@ -0,0 +1,23 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""The vertexai module."""
+
+from . import evals
+from .client import Client
+
+__all__ = [
+    "Client",
+    "evals",
+]
diff --git a/vertexai/_genai/client.py b/vertexai/_genai/client.py
@@ -0,0 +1,92 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from typing import Optional, Union
+
+import google.auth
+from google.genai import client
+from google.genai import types
+from .evals import Evals
+from .evals import AsyncEvals
+
+
+class AsyncClient:
+    """Async Client for the GenAI SDK."""
+
+    def __init__(self, api_client: client.Client):
+        self._api_client = api_client
+        self._aio = AsyncClient(self._api_client)
+        self._evals = AsyncEvals(self._api_client)
+
+    @property
+    def evals(self) -> AsyncEvals:
+        return self._evals
+
+
+class Client:
+    """Client for the GenAI SDK.
+
+    Use this client to interact with Vertex-specific Gemini features.
+    """
+
+    def __init__(
+        self,
+        *,
+        credentials: Optional[google.auth.credentials.Credentials] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        debug_config: Optional[client.DebugConfig] = None,
+        http_options: Optional[Union[types.HttpOptions, types.HttpOptionsDict]] = None,
+    ):
+        """Initializes the client.
+
+        Args:
+           credentials (google.auth.credentials.Credentials): The credentials to use
+             for authentication when calling the Vertex AI APIs. Credentials can be
+             obtained from environment variables and default credentials. For more
+             information, see `Set up Application Default Credentials
+             <https://cloud.google.com/docs/authentication/provide-credentials-adc>`_.
+           project (str): The `Google Cloud project ID
+             <https://cloud.google.com/vertex-ai/docs/start/cloud-environment>`_ to
+             use for quota. Can be obtained from environment variables (for example,
+             ``GOOGLE_CLOUD_PROJECT``).
+           location (str): The `location
+             <https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations>`_
+             to send API requests to (for example, ``us-central1``). Can be obtained
+             from environment variables.
+           debug_config (DebugConfig): Config settings that control network behavior
+             of the client. This is typically used when running test code.
+           http_options (Union[HttpOptions, HttpOptionsDict]): Http options to use
+             for the client.
+        """
+
+        self._debug_config = debug_config or client.DebugConfig()
+        if isinstance(http_options, dict):
+            http_options = types.HttpOptions(**http_options)
+
+        self._api_client = client.Client._get_api_client(
+            vertexai=True,
+            credentials=credentials,
+            project=project,
+            location=location,
+            debug_config=self._debug_config,
+            http_options=http_options,
+        )
+
+        self._evals = Evals(self._api_client)
+
+    @property
+    def evals(self) -> Evals:
+        return self._evals
diff --git a/vertexai/_genai/evals.py b/vertexai/_genai/evals.py
diff --git a/vertexai/_genai/types.py b/vertexai/_genai/types.py
