feat: Eval SDK: Migrate model call method by genai SDK usage

PiperOrigin-RevId: 893061063

Author: vertex-sdk-bot

tests/unit/vertexai/test_evaluation.py

@@ -22,6 +22,7 @@
 from unittest import mock
 
 from google import auth
+from google import genai
 from google.auth import credentials as auth_credentials
 from google.cloud import aiplatform
 import vertexai
@@ -1025,6 +1026,62 @@ def test_compute_pointwise_metrics_metric_prompt_template_example(
             "explanation",
         ]
 
+    @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
+    def test_compute_pointwise_metrics_metric_prompt_template_example_string_model(
+        self, api_transport
+    ):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            api_transport=api_transport,
+        )
+        mock_client = mock.create_autospec(genai.Client, instance=True)
+        mock_response = mock.MagicMock()
+        mock_response.text = "test_response"
+        mock_client.models.generate_content.return_value = mock_response
+
+        test_metrics = [Pointwise.SUMMARIZATION_QUALITY]
+        test_eval_task = EvalTask(
+            dataset=_TEST_EVAL_DATASET_WITHOUT_RESPONSE, metrics=test_metrics
+        )
+        mock_metric_results = _MOCK_SUMMARIZATION_QUALITY_RESULT
+        with mock.patch.object(genai, "Client", return_value=mock_client):
+            with mock.patch.object(
+                target=gapic_evaluation_services.EvaluationServiceClient,
+                attribute="evaluate_instances",
+                side_effect=mock_metric_results,
+            ):
+                test_result = test_eval_task.evaluate(
+                    model="gemini-1.5-pro",
+                    prompt_template="{instruction} test prompt template {context}",
+                )
+
+        assert test_result.summary_metrics["row_count"] == 2
+        assert test_result.summary_metrics["summarization_quality/mean"] == 4.5
+        assert test_result.summary_metrics[
+            "summarization_quality/std"
+        ] == pytest.approx(0.7, 0.1)
+        assert set(test_result.metrics_table.columns.values) == set(
+            [
+                "context",
+                "instruction",
+                "reference",
+                "prompt",
+                "response",
+                "summarization_quality/score",
+                "summarization_quality/explanation",
+            ]
+        )
+        assert list(
+            test_result.metrics_table["summarization_quality/score"].values
+        ) == [5, 4]
+        assert list(
+            test_result.metrics_table["summarization_quality/explanation"].values
+        ) == [
+            "explanation",
+            "explanation",
+        ]
+
     @pytest.mark.parametrize("api_transport", ["grpc", "rest"])
     def test_compute_pointwise_metrics_without_model_inference(self, api_transport):
         aiplatform.init(
@@ -1401,13 +1458,13 @@ def test_compute_multiple_metrics(self, api_transport):
         mock_baseline_model.generate_content.return_value = (
             _MOCK_MODEL_INFERENCE_RESPONSE
         )
-        mock_baseline_model._model_name = "publishers/google/model/gemini-pro"
+        mock_baseline_model._model_name = "gemini-2.5-pro"
         _TEST_PAIRWISE_METRIC._baseline_model = mock_baseline_model
         mock_model = mock.create_autospec(
             generative_models.GenerativeModel, instance=True
         )
         mock_model.generate_content.return_value = _MOCK_MODEL_INFERENCE_RESPONSE
-        mock_model._model_name = "publishers/google/model/gemini-pro"
+        mock_model._model_name = "gemini-2.5-flash"
         test_metrics = [
             "exact_match",
             Pointwise.SUMMARIZATION_QUALITY,
@@ -2654,6 +2711,31 @@ def test_default_rubrics_parser_with_invalid_json(self):
         parsed_rubrics = utils_preview.parse_rubrics(_INVALID_UNPARSED_RUBRIC)
         assert parsed_rubrics == {"questions": ""}
 
+    def test_generate_responses_from_genai_model(self):
+        mock_client = mock.create_autospec(genai.Client, instance=True)
+        mock_response = mock.MagicMock()
+        mock_response.text = "test_response"
+        mock_client.models.generate_content.return_value = mock_response
+
+        with mock.patch.object(genai, "Client", return_value=mock_client):
+            evaluation_run_config = eval_base.EvaluationRunConfig(
+                dataset=_TEST_EVAL_DATASET_WITHOUT_RESPONSE.copy(),
+                metrics=[],
+                metric_column_mapping={},
+                client=mock.MagicMock(),
+                evaluation_service_qps=1,
+                retry_timeout=1,
+            )
+            _evaluation._generate_responses_from_genai_model(
+                "gemini-2.5-pro", evaluation_run_config
+            )
+
+        assert list(evaluation_run_config.dataset["response"].values) == [
+            "test_response",
+            "test_response",
+        ]
+        assert mock_client.models.generate_content.call_count == 2
+
     def test_generate_responses_from_gemini_model(self):
         mock_model = mock.create_autospec(
             generative_models.GenerativeModel, instance=True

vertexai/evaluation/_evaluation.py

@@ -22,6 +22,8 @@
 import time
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TYPE_CHECKING, Union
 
+from google import genai
+from google.cloud import aiplatform
 from google.cloud.aiplatform import base
 from google.cloud.aiplatform_v1beta1.types import (
     content as gapic_content_types,
@@ -373,6 +375,106 @@ def _generate_content_text_response(
     return constants.RESPONSE_ERROR
 
 
+def _generate_content_text_response_genai(
+    model: str, client: genai.Client, prompt: str, max_retries: int = 3
+) -> str:
+    """Generates a text response from Gemini model from a text prompt with retries using genai module.
+
+    Args:
+        model: The model name string.
+        client: The genai client instance.
+        prompt: The prompt to send to the model.
+        max_retries: Maximum number of retries for response generation.
+
+    Returns:
+        The text response from the model.
+        Returns constants.RESPONSE_ERROR if there is an error after all retries.
+    """
+    for retry_attempt in range(max_retries):
+        try:
+            response = client.models.generate_content(
+                model=model,
+                contents=prompt,
+            )
+            # The new SDK raises exceptions on blocked content instead of returning
+            # block_reason directly, so if it succeeds, we can return the text.
+            if response.text:
+                return response.text
+            else:
+                _LOGGER.warning(
+                    "The model response was empty or blocked.\n"
+                    f"Prompt: {prompt}.\n"
+                    f"Retry attempt: {retry_attempt + 1}/{max_retries}"
+                )
+        except Exception as e:  # pylint: disable=broad-except
+            error_message = (
+                f"Failed to generate response candidates from GenAI model "
+                f"{model}.\n"
+                f"Error: {e}.\n"
+                f"Prompt: {prompt}.\n"
+                f"Retry attempt: {retry_attempt + 1}/{max_retries}"
+            )
+            _LOGGER.warning(error_message)
+        if retry_attempt < max_retries - 1:
+            _LOGGER.info(
+                f"Retrying response generation for prompt: {prompt}, attempt "
+                f"{retry_attempt + 1}/{max_retries}..."
+            )
+
+    final_error_message = (
+        f"Failed to generate response from GenAI model {model}.\n" f"Prompt: {prompt}."
+    )
+    _LOGGER.warning(final_error_message)
+    return constants.RESPONSE_ERROR
+
+
+def _generate_responses_from_genai_model(
+    model: str,
+    evaluation_run_config: evaluation_base.EvaluationRunConfig,
+    is_baseline_model: bool = False,
+) -> None:
+    """Generates responses from Gemini model using genai module.
+
+    Args:
+        model: The model name string.
+        evaluation_run_config: Evaluation Run Configurations.
+        is_baseline_model: Whether the model is a baseline model for PairwiseMetric.
+    """
+    df = evaluation_run_config.dataset.copy()
+
+    _LOGGER.info(
+        f"Generating a total of {evaluation_run_config.dataset.shape[0]} "
+        f"responses from Gemini model {model} using genai module."
+    )
+    tasks = []
+    client = genai.Client(
+        vertexai=True,
+        project=aiplatform.initializer.global_config.project,
+        location=aiplatform.initializer.global_config.location,
+    )
+    with tqdm(total=len(df)) as pbar:
+        with futures.ThreadPoolExecutor(max_workers=constants.MAX_WORKERS) as executor:
+            for _, row in df.iterrows():
+                task = executor.submit(
+                    _generate_content_text_response_genai,
+                    prompt=row[constants.Dataset.PROMPT_COLUMN],
+                    model=model,
+                    client=client,
+                )
+                task.add_done_callback(lambda _: pbar.update(1))
+                tasks.append(task)
+        responses = [future.result() for future in tasks]
+    if is_baseline_model:
+        evaluation_run_config.dataset = df.assign(baseline_model_response=responses)
+    else:
+        evaluation_run_config.dataset = df.assign(response=responses)
+
+    _LOGGER.info(
+        f"All {evaluation_run_config.dataset.shape[0]} responses are successfully "
+        f"generated from Gemini model {model} using genai module."
+    )
+
+
 def _generate_responses_from_gemini_model(
     model: generative_models.GenerativeModel,
     evaluation_run_config: evaluation_base.EvaluationRunConfig,
@@ -463,7 +565,7 @@ def _generate_response_from_custom_model_fn(
 
 
 def _run_model_inference(
-    model: Union[generative_models.GenerativeModel, Callable[[str], str]],
+    model: Union[str, generative_models.GenerativeModel, Callable[[str], str]],
     evaluation_run_config: evaluation_base.EvaluationRunConfig,
     response_column_name: str = constants.Dataset.MODEL_RESPONSE_COLUMN,
 ) -> None:
@@ -488,9 +590,18 @@ def _run_model_inference(
             if constants.Dataset.PROMPT_COLUMN in evaluation_run_config.dataset.columns:
                 t1 = time.perf_counter()
                 if isinstance(model, generative_models.GenerativeModel):
+                    _LOGGER.warning(
+                        "vertexai.generative_models.GenerativeModel is deprecated for "
+                        "evaluation and will be removed in June 2026. Please pass a "
+                        "string model name instead."
+                    )
                     _generate_responses_from_gemini_model(
                         model, evaluation_run_config, is_baseline_model
                     )
+                elif isinstance(model, str):
+                    _generate_responses_from_genai_model(
+                        model, evaluation_run_config, is_baseline_model
+                    )
                 elif callable(model):
                     _generate_response_from_custom_model_fn(
                         model, evaluation_run_config, is_baseline_model
@@ -878,7 +989,7 @@ def evaluate(
     metrics: List[Union[str, metrics_base._Metric]],
     *,
     model: Optional[
-        Union[generative_models.GenerativeModel, Callable[[str], str]]
+        Union[str, generative_models.GenerativeModel, Callable[[str], str]]
     ] = None,
     prompt_template: Optional[Union[str, prompt_template_base.PromptTemplate]] = None,
     metric_column_mapping: Dict[str, str],