feat: Add GenAI client (experimental)

PiperOrigin-RevId: 766750716

Author: sararob

noxfile.py

@@ -67,6 +67,8 @@
     "pytest-asyncio",
     # Preventing: py.test: error: unrecognized arguments: -n=auto --dist=loadscope
     "pytest-xdist",
+    # "pandas",
+    # "tqdm",
 ]
 UNIT_TEST_EXTERNAL_DEPENDENCIES = []
 UNIT_TEST_LOCAL_DEPENDENCIES = []

tests/unit/architecture/test_vertexai_import.py

@@ -90,6 +90,28 @@ def test_vertexai_import():
     new_modules = modules_after_vertexai - modules_before_aip
     _test_external_imports(new_modules)
 
+    # Tests the GenAI client module is lazy loaded.
+    from vertexai._genai import client as _  # noqa: F401,F811
+
+    modules_after_genai_client_import = set(sys.modules)
+
+    # The evals module has additional required deps that should not be required
+    # to instantiate a client.
+    assert "pandas" not in modules_after_genai_client_import
+
+    # Tests the evals module is lazy loaded.
+    eval_module_deps_subset = ("pydantic", "mcp", "pandas")
+
+    assert not all(module in new_modules for module in eval_module_deps_subset)
+
+    from vertexai._genai import evals as _  # noqa: F401,F811
+
+    modules_after_eval_import = set(sys.modules)
+
+    assert all(
+        module in modules_after_eval_import for module in eval_module_deps_subset
+    )
+
 
 def _test_external_imports(new_modules: list):
     builtin_modules = {

tests/unit/vertexai/genai/test_evals.py

@@ -18,22 +18,29 @@
 import json
 import os
 from unittest import mock
+import warnings
 
 from google.cloud import aiplatform
 import vertexai
 from google.cloud.aiplatform import initializer as aiplatform_initializer
 from vertexai import _genai
+from vertexai._genai import evals
 from vertexai._genai import types as vertexai_genai_types
+from google.genai import client
+from google.genai import errors as genai_errors
 from google.genai import types as genai_types
-import google.genai.errors as genai_errors
 import pandas as pd
 import pytest
-import warnings
+
 
 _TEST_PROJECT = "test-project"
 _TEST_LOCATION = "us-central1"
 
 
+_genai.evals._lazy_load_evals_common()
+_evals_common = _genai.evals._evals_common
+_evals_utils = _genai._evals_utils
+
 pytestmark = pytest.mark.usefixtures("google_auth_mock")
 
 
@@ -48,47 +55,41 @@ def setup_method(self):
             project=_TEST_PROJECT,
             location=_TEST_LOCATION,
         )
+        self.client = vertexai.Client(project=_TEST_PROJECT, 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
-        )
+    @mock.patch.object(client.Client, "_get_api_client")
+    @mock.patch.object(evals.Evals, "_evaluate_instances")
+    def test_evaluate_instances(self, mock_evaluate, mock_get_api_client):
         with warnings.catch_warnings(record=True) as captured_warnings:
             warnings.simplefilter("always")
-            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()
-                )
-                assert captured_warnings[0].category == genai_errors.ExperimentalWarning
+            self.client.evals._evaluate_instances(
+                bleu_input=vertexai_genai_types.BleuInput()
+            )
+            mock_evaluate.assert_called_once_with(
+                bleu_input=vertexai_genai_types.BleuInput()
+            )
+            assert captured_warnings[0].category == genai_errors.ExperimentalWarning
 
     @pytest.mark.usefixtures("google_auth_mock")
     def test_eval_run(self):
-        test_client = _genai.client.Client(
-            project=_TEST_PROJECT, location=_TEST_LOCATION
-        )
+        test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
         with pytest.raises(NotImplementedError):
             test_client.evals.run()
 
     @pytest.mark.usefixtures("google_auth_mock")
-    def test_eval_batch_eval(self):
-        test_client = _genai.client.Client(
-            project=_TEST_PROJECT, location=_TEST_LOCATION
-        )
-        with mock.patch.object(test_client.evals, "batch_eval") as mock_batch_eval:
-            test_client.evals.batch_eval(
-                dataset=_genai.types.EvaluationDataset(),
-                metrics=[_genai.types.Metric()],
-                output_config=_genai.types.OutputConfig(),
-                autorater_config=_genai.types.AutoraterConfig(),
-                config=_genai.types.EvaluateDatasetConfig(),
-            )
-            mock_batch_eval.assert_called_once()
+    @mock.patch.object(client.Client, "_get_api_client")
+    @mock.patch.object(evals.Evals, "batch_eval")
+    def test_eval_batch_eval(self, mock_evaluate, mock_get_api_client):
+        test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
+        test_client.evals.batch_eval(
+            dataset=vertexai_genai_types.EvaluationDataset(),
+            metrics=[vertexai_genai_types.Metric()],
+            output_config=vertexai_genai_types.OutputConfig(),
+            autorater_config=vertexai_genai_types.AutoraterConfig(),
+            config=vertexai_genai_types.EvaluateDatasetConfig(),
+        )
+        mock_evaluate.assert_called_once()
 
 
 class TestEvalsClientInference:
@@ -99,20 +100,16 @@ def setup_method(self):
         importlib.reload(aiplatform)
         importlib.reload(vertexai)
         importlib.reload(_genai.client)
-        importlib.reload(_genai.types)
+        importlib.reload(vertexai_genai_types)
         importlib.reload(_genai.evals)
-        importlib.reload(_genai._evals_utils)
-        importlib.reload(_genai._evals_common)
         vertexai.init(
             project=_TEST_PROJECT,
             location=_TEST_LOCATION,
         )
-        self.client = _genai.client.Client(
-            project=_TEST_PROJECT, location=_TEST_LOCATION
-        )
+        self.client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch.object(_evals_common, "Models")
+    @mock.patch.object(_evals_utils, "EvalDatasetLoader")
     def test_inference_with_string_model_success(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -153,7 +150,7 @@ def test_inference_with_string_model_success(
             ),
         )
 
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch.object(_evals_utils, "EvalDatasetLoader")
     def test_inference_with_callable_model_success(self, mock_eval_dataset_loader):
         mock_df = pd.DataFrame({"prompt": ["test prompt"]})
         mock_eval_dataset_loader.return_value.load.return_value = mock_df.to_dict(
@@ -178,8 +175,8 @@ def mock_model_fn(contents):
             ),
         )
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch.object(_evals_common, "Models")
+    @mock.patch.object(_evals_utils, "EvalDatasetLoader")
     def test_inference_with_prompt_template(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -202,7 +199,7 @@ def test_inference_with_prompt_template(
             mock_generate_content_response
         )
 
-        config = _genai.types.EvalRunInferenceConfig(
+        config = vertexai_genai_types.EvalRunInferenceConfig(
             prompt_template="Hello {text_input}"
         )
         result_df = self.client.evals.run_inference(
@@ -223,9 +220,9 @@ def test_inference_with_prompt_template(
             ),
         )
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
-    @mock.patch(f"{_genai._evals_utils.__name__}.GcsUtils")
+    @mock.patch.object(_evals_common, "Models")
+    @mock.patch.object(_evals_utils, "EvalDatasetLoader")
+    @mock.patch.object(_evals_utils, "GcsUtils")
     def test_inference_with_gcs_destination(
         self, mock_gcs_utils, mock_eval_dataset_loader, mock_models
     ):
@@ -249,7 +246,7 @@ def test_inference_with_gcs_destination(
         )
 
         gcs_dest_path = "gs://bucket/output.jsonl"
-        config = _genai.types.EvalRunInferenceConfig(dest=gcs_dest_path)
+        config = vertexai_genai_types.EvalRunInferenceConfig(dest=gcs_dest_path)
 
         result_df = self.client.evals.run_inference(
             model="gemini-pro", src=mock_df, config=config
@@ -270,8 +267,8 @@ def test_inference_with_gcs_destination(
         )
         pd.testing.assert_frame_equal(result_df, expected_df_to_save)
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch.object(_evals_common, "Models")
+    @mock.patch.object(_evals_utils, "EvalDatasetLoader")
     @mock.patch("pandas.DataFrame.to_json")
     @mock.patch("os.makedirs")
     def test_inference_with_local_destination(
@@ -301,7 +298,7 @@ def test_inference_with_local_destination(
         )
 
         local_dest_path = "/tmp/test/output_dir/results.jsonl"
-        config = _genai.types.EvalRunInferenceConfig(dest=local_dest_path)
+        config = vertexai_genai_types.EvalRunInferenceConfig(dest=local_dest_path)
 
         result_df = self.client.evals.run_inference(
             model="gemini-pro", src=mock_df, config=config
@@ -319,8 +316,8 @@ def test_inference_with_local_destination(
         )
         pd.testing.assert_frame_equal(result_df, expected_df)
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch.object(_evals_common, "Models")
+    @mock.patch.object(_evals_utils, "EvalDatasetLoader")
     def test_inference_from_request_column_save_locally(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -359,7 +356,7 @@ def test_inference_from_request_column_save_locally(
         )
 
         local_dest_path = "/tmp/output.jsonl"
-        config = _genai.types.EvalRunInferenceConfig(dest=local_dest_path)
+        config = vertexai_genai_types.EvalRunInferenceConfig(dest=local_dest_path)
 
         result_df = self.client.evals.run_inference(
             model="gemini-pro", src=mock_df, config=config
@@ -396,7 +393,7 @@ def test_inference_from_request_column_save_locally(
         assert saved_records == expected_records
         os.remove(local_dest_path)
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
+    @mock.patch.object(_evals_common, "Models")
     def test_inference_from_local_jsonl_file(self, mock_models):
         # Create a temporary JSONL file
         local_src_path = "/tmp/input.jsonl"
@@ -450,7 +447,7 @@ def test_inference_from_local_jsonl_file(self, mock_models):
         pd.testing.assert_frame_equal(result_df, expected_df)
         os.remove(local_src_path)
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
+    @mock.patch.object(_evals_common, "Models")
     def test_inference_from_local_csv_file(self, mock_models):
         # Create a temporary CSV file
         local_src_path = "/tmp/input.csv"
@@ -501,8 +498,8 @@ def test_inference_from_local_csv_file(self, mock_models):
         pd.testing.assert_frame_equal(result_df, expected_df)
         os.remove(local_src_path)
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch.object(_evals_common, "Models")
+    @mock.patch.object(_evals_utils, "EvalDatasetLoader")
     def test_inference_with_row_level_config_overrides(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -584,8 +581,8 @@ def test_inference_with_row_level_config_overrides(
         )
         pd.testing.assert_frame_equal(result_df, expected_df)
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch.object(_evals_common, "Models")
+    @mock.patch.object(_evals_utils, "EvalDatasetLoader")
     def test_inference_with_multimodal_content(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -623,7 +620,7 @@ def test_inference_with_multimodal_content(
             mock_generate_content_response
         )
 
-        config = _genai.types.EvalRunInferenceConfig(
+        config = vertexai_genai_types.EvalRunInferenceConfig(
             prompt_template="multimodal prompt: {media_content}{text_input}"
         )
         result_df = self.client.evals.run_inference(

tests/unit/vertexai/genai/test_genai_client.py

@@ -20,7 +20,6 @@
 from google.cloud import aiplatform
 import vertexai
 from google.cloud.aiplatform import initializer as aiplatform_initializer
-from vertexai import _genai
 import pytest
 
 _TEST_PROJECT = "test-project"
@@ -44,9 +43,7 @@ def setup_method(self):
 
     @pytest.mark.usefixtures("google_auth_mock")
     def test_genai_client(self):
-        test_client = _genai.client.Client(
-            project=_TEST_PROJECT, location=_TEST_LOCATION
-        )
+        test_client = vertexai.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

vertexai/__init__.py

@@ -22,10 +22,8 @@
 
 from google.cloud.aiplatform import init
 
-__all__ = [
-    "init",
-    "preview",
-]
+_genai_client = None
+_genai_types = None
 
 
 def __getattr__(name):
@@ -39,4 +37,24 @@ def __getattr__(name):
         # `import google.cloud.aiplatform.vertexai.preview as vertexai_preview`
         return importlib.import_module(".preview", __name__)
 
+    if name == "Client":
+        global _genai_client
+        if _genai_client is None:
+            _genai_client = importlib.import_module("._genai.client", __name__)
+        return getattr(_genai_client, name)
+
+    if name == "types":
+        global _genai_types
+        if _genai_types is None:
+            _genai_types = importlib.import_module("._genai.types", __name__)
+        return _genai_types
+
     raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+__all__ = [
+    "init",
+    "preview",
+    "Client",
+    "types",
+]

vertexai/_genai/__init__.py

@@ -14,9 +14,29 @@
 #
 """The vertexai module."""
 
-from . import evals
+import importlib
+
 from .client import Client
 
+_evals = None
+
+
+def __getattr__(name):
+    if name == "evals":
+        global _evals
+        if _evals is None:
+            try:
+                _evals = importlib.import_module(".evals", __package__)
+            except ImportError as e:
+                raise ImportError(
+                    "The 'evals' module requires 'pandas' and 'tqdm'. "
+                    "Please install them using pip install "
+                    "google-cloud-aiplatform[evaluation]"
+                ) from e
+        return _evals
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
 __all__ = [
     "Client",
     "evals",