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/vertexai/genai/test_evals.py

@@ -25,7 +25,8 @@
 from vertexai import _genai
 from vertexai._genai import types as vertexai_genai_types
 from google.genai import types as genai_types
-import google.genai.errors as genai_errors
+from google.genai import client
+from vertexai._genai import evals
 import pandas as pd
 import pytest
 import warnings
@@ -48,47 +49,35 @@ 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
-        )
-        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
+    @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):
+
+        self.client.evals._evaluate_instances(bleu_input=vertexai_genai_types.BleuInput())
+        mock_evaluate.assert_called_once_with(bleu_input=vertexai_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
-        )
+        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
+    @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(),
         )
-        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_evaluate.assert_called_once()
 
 
 class TestEvalsClientInference:
@@ -99,20 +88,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("google.cloud.aiplatform.vertexai._genai._evals_common.Models")
+    @mock.patch("google.cloud.aiplatform.vertexai._genai._evals_utils.EvalDatasetLoader")
     def test_inference_with_string_model_success(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -153,7 +138,7 @@ def test_inference_with_string_model_success(
             ),
         )
 
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._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 +163,8 @@ def mock_model_fn(contents):
             ),
         )
 
-    @mock.patch(f"{_genai._evals_common.__name__}.Models")
-    @mock.patch(f"{_genai._evals_utils.__name__}.EvalDatasetLoader")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._evals_common.Models")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._evals_utils.EvalDatasetLoader")
     def test_inference_with_prompt_template(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -202,7 +187,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 +208,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(f"google.cloud.aiplatform.vertexai._genai._evals_common.Models")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._evals_utils.EvalDatasetLoader")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._evals_utils.GcsUtils")
     def test_inference_with_gcs_destination(
         self, mock_gcs_utils, mock_eval_dataset_loader, mock_models
     ):
@@ -249,7 +234,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 +255,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(f"google.cloud.aiplatform.vertexai._genai._evals_common.Models")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._evals_utils.EvalDatasetLoader")
     @mock.patch("pandas.DataFrame.to_json")
     @mock.patch("os.makedirs")
     def test_inference_with_local_destination(
@@ -301,7 +286,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 +304,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(f"google.cloud.aiplatform.vertexai._genai._evals_common.Models")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._evals_utils.EvalDatasetLoader")
     def test_inference_from_request_column_save_locally(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -359,7 +344,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 +381,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(f"google.cloud.aiplatform.vertexai._genai._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 +435,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(f"google.cloud.aiplatform.vertexai._genai._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 +486,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(f"google.cloud.aiplatform.vertexai._genai._evals_common.Models")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._evals_utils.EvalDatasetLoader")
     def test_inference_with_row_level_config_overrides(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -584,8 +569,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(f"google.cloud.aiplatform.vertexai._genai._evals_common.Models")
+    @mock.patch(f"google.cloud.aiplatform.vertexai._genai._evals_utils.EvalDatasetLoader")
     def test_inference_with_multimodal_content(
         self, mock_eval_dataset_loader, mock_models
     ):
@@ -623,7 +608,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

@@ -44,9 +44,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,10 @@
 
 from google.cloud.aiplatform import init
 
-__all__ = [
-    "init",
-    "preview",
-]
+# from vertexai._genai.client import Client
+# from vertexai._genai import types
+
+_genai_client = None
 
 
 def __getattr__(name):
@@ -39,4 +39,18 @@ def __getattr__(name):
         # `import google.cloud.aiplatform.vertexai.preview as vertexai_preview`
         return importlib.import_module(".preview", __name__)
 
+    if name in ("Client", "types"):
+        global _genai_client
+        if _genai_client is None:
+            _genai_client = importlib.import_module("._genai.client", __name__)
+        return getattr(_genai_client, name)
+
     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
+_evals_common = 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(
+                    f"The 'evals' module requires 'pandas' and 'tqdm'. "
+                    f"Please install them using pip install "
+                    f"google-cloud-aiplatform[evaluation]"
+                ) from e
+        return _evals
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
 __all__ = [
     "Client",
     "evals",

vertexai/_genai/_evals_common.py

@@ -24,12 +24,26 @@
 from google.genai import types as genai_types
 from google.genai._api_client import BaseApiClient
 from google.genai.models import Models
-import pandas as pd
-from tqdm import tqdm
 
 from . import _evals_utils
 from . import types
 
+try:
+    import pandas as pd
+except ImportError:
+    raise ImportError(
+        'Pandas is not installed. Please install the SDK using "pip install'
+        ' google-cloud-aiplatform[evaluation]"'
+    )
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    raise ImportError(
+        'tqdm is not installed. Please install the SDK using "pip install'
+        ' google-cloud-aiplatform[evaluation]"'
+    )
+
 logger = logging.getLogger(__name__)
 
 _MAX_WORKERS = 100

vertexai/_genai/_evals_utils.py

@@ -22,7 +22,13 @@
 from google.cloud import bigquery
 from google.cloud import storage
 from google.genai._api_client import BaseApiClient
-import pandas as pd
+try:
+    import pandas as pd
+except ImportError:
+    raise ImportError(
+        'Pandas is not installed. Please install the SDK using "pip install'
+        ' google-cloud-aiplatform[evaluation]"'
+    )
 
 logger = logging.getLogger(__name__)