feat: GenAI Client(evals) - support setting autorater generation config for predefined rubric metrics

PiperOrigin-RevId: 833487047

Author: jsondai

tests/unit/vertexai/genai/replays/test_evaluate_predefined_metrics.py

@@ -16,6 +16,7 @@
 
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai import types
+from google.genai import types as genai_types
 import pandas as pd
 
 
@@ -60,6 +61,50 @@ def test_evaluation_result(client):
         assert case_result.response_candidate_results is not None
 
 
+def test_evaluation_result_with_autorater_config(client):
+    """Tests that evaluate() produces a correctly structured EvaluationResult."""
+    prompts_df = pd.DataFrame(
+        {
+            "prompt": ["Explain the concept of machine learning in simple terms."],
+            "response": [
+                "Machine learning is a type of artificial intelligence that allows"
+                " computers to learn from data without being explicitly programmed."
+            ],
+        }
+    )
+
+    eval_dataset = types.EvaluationDataset(
+        eval_dataset_df=prompts_df,
+        candidate_name="gemini-2.5-flash",
+    )
+
+    predefined_metric_with_autorater_config = types.RubricMetric.GENERAL_QUALITY(
+        judge_model_generation_config=genai_types.GenerationConfig(
+            temperature=0.1,
+            max_output_tokens=1024,
+        )
+    )
+
+    evaluation_result = client.evals.evaluate(
+        dataset=eval_dataset,
+        metrics=[predefined_metric_with_autorater_config],
+    )
+
+    assert isinstance(evaluation_result, types.EvaluationResult)
+
+    assert evaluation_result.summary_metrics is not None
+    for summary in evaluation_result.summary_metrics:
+        assert isinstance(summary, types.AggregatedMetricResult)
+        assert summary.metric_name == "general_quality_v1"
+        assert summary.mean_score is not None
+
+    assert evaluation_result.eval_case_results is not None
+    for case_result in evaluation_result.eval_case_results:
+        assert isinstance(case_result, types.EvalCaseResult)
+        assert case_result.eval_case_index is not None
+        assert case_result.response_candidate_results is not None
+
+
 def test_multi_turn_predefined_metric(client):
     """Tests that evaluate works with multi-turn predefined metrics."""
     prompts_data = {

tests/unit/vertexai/genai/test_evals.py

@@ -189,6 +189,51 @@ def test_eval_evaluate_with_agent_info(self, mock_execute_evaluation):
         assert "agent_info" in kwargs
         assert kwargs["agent_info"] == agent_info
 
+    def test_evaluate_predefined_metric_with_autorater_config(self):
+        dataset = vertexai_genai_types.EvaluationDataset(
+            eval_dataset_df=pd.DataFrame([{"prompt": "p1", "response": "r1"}])
+        )
+        generation_config = genai_types.GenerationConfig(
+            temperature=0.1,
+            max_output_tokens=1024,
+        )
+        metrics = [
+            vertexai_genai_types.RubricMetric.GENERAL_QUALITY(
+                judge_model_generation_config=generation_config
+            )
+        ]
+
+        mock_prebuilt_metric = vertexai_genai_types.RubricMetric(
+            name="general_quality_v1",
+            judge_model_generation_config=generation_config,
+        )
+        mock_prebuilt_metric._is_predefined = True
+        mock_prebuilt_metric._version = "v1"
+
+        with mock.patch(
+            "vertexai._genai._evals_metric_loaders.LazyLoadedPrebuiltMetric._fetch_and_parse",
+            return_value=mock_prebuilt_metric,
+        ), mock.patch(
+            "vertexai._genai.evals.Evals._evaluate_instances"
+        ) as mock_evaluate_instances, mock.patch(
+            "vertexai._genai._evals_metric_handlers._evals_constant.SUPPORTED_PREDEFINED_METRICS",
+            frozenset(["general_quality_v1"]),
+        ):
+            mock_evaluate_instances.return_value = (
+                vertexai_genai_types.EvaluateInstancesResponse(
+                    metric_results=[vertexai_genai_types.MetricResult(score=0.9)]
+                )
+            )
+            self.client.evals.evaluate(
+                dataset=dataset,
+                metrics=metrics,
+            )
+
+            mock_evaluate_instances.assert_called_once()
+            _, kwargs = mock_evaluate_instances.call_args
+            assert "autorater_config" in kwargs
+            assert kwargs["autorater_config"].generation_config == generation_config
+
 
 class TestEvalsVisualization:
     @mock.patch(
@@ -4990,7 +5035,9 @@ def test_execute_evaluation_adds_creation_timestamp(
         frozenset(["summarization_quality"]),
     )
     @mock.patch("time.sleep", return_value=None)
-    @mock.patch("vertexai._genai.evals.Evals._evaluate_instances")
+    @mock.patch(
+        "vertexai._genai.evals.Evals._evaluate_instances"
+    )
     def test_predefined_metric_retry_on_resource_exhausted(
         self,
         mock_private_evaluate_instances,
@@ -5043,7 +5090,9 @@ def test_predefined_metric_retry_on_resource_exhausted(
         frozenset(["summarization_quality"]),
     )
     @mock.patch("time.sleep", return_value=None)
-    @mock.patch("vertexai._genai.evals.Evals._evaluate_instances")
+    @mock.patch(
+        "vertexai._genai.evals.Evals._evaluate_instances"
+    )
     def test_predefined_metric_retry_fail_on_resource_exhausted(
         self,
         mock_private_evaluate_instances,

vertexai/_genai/_evals_metric_handlers.py

@@ -620,6 +620,10 @@ def _add_autorater_config(self, payload: dict[str, Any]) -> None:
         autorater_config = {}
         if self.metric.judge_model:
             autorater_config["autorater_model"] = self.metric.judge_model
+        if self.metric.judge_model_generation_config:
+            autorater_config["generation_config"] = (
+                self.metric.judge_model_generation_config
+            )
         if self.metric.judge_model_sampling_count:
             autorater_config["sampling_count"] = self.metric.judge_model_sampling_count  # type: ignore[assignment]
 
@@ -986,10 +990,25 @@ def _build_request_payload(
             agent_data=PredefinedMetricHandler._eval_case_to_agent_data(eval_case),
         )
 
-        return {
+        request_payload = {
             "instance": instance_payload,
         }
 
+        autorater_config = {}
+        if self.metric.judge_model:
+            autorater_config["autorater_model"] = self.metric.judge_model
+        if self.metric.judge_model_generation_config:
+            autorater_config["generation_config"] = (
+                self.metric.judge_model_generation_config
+            )
+        if self.metric.judge_model_sampling_count:
+            autorater_config["sampling_count"] = self.metric.judge_model_sampling_count
+        if autorater_config:
+            request_payload["autorater_config"] = genai_types.AutoraterConfig(
+                **autorater_config
+            )
+        return request_payload
+
     @override
     def get_metric_result(
         self, eval_case: types.EvalCase, response_index: int
@@ -1001,7 +1020,9 @@ def get_metric_result(
             for attempt in range(_MAX_RETRIES):
                 try:
                     api_response = self.module._evaluate_instances(
-                        metrics=[self.metric], instance=payload.get("instance")
+                        metrics=[self.metric],
+                        instance=payload.get("instance"),
+                        autorater_config=payload.get("autorater_config"),
                     )
                     break
                 except genai_errors.ClientError as e:

vertexai/_genai/types/common.py

@@ -2622,6 +2622,10 @@ class Metric(_common.BaseModel):
     judge_model: Optional[str] = Field(
         default=None, description="""The judge model for the metric."""
     )
+    judge_model_generation_config: Optional[genai_types.GenerationConfig] = Field(
+        default=None,
+        description="""The generation config for the judge LLM (temperature, top_k, top_p, etc).""",
+    )
     judge_model_sampling_count: Optional[int] = Field(
         default=None, description="""The sampling count for the judge model."""
     )
@@ -2825,6 +2829,9 @@ class MetricDict(TypedDict, total=False):
     judge_model: Optional[str]
     """The judge model for the metric."""
 
+    judge_model_generation_config: Optional[genai_types.GenerationConfigDict]
+    """The generation config for the judge LLM (temperature, top_k, top_p, etc)."""
+
     judge_model_sampling_count: Optional[int]
     """The sampling count for the judge model."""