feat: Allow using registered metric resource names in evaluation

PiperOrigin-RevId: 880868820

Author: vertex-sdk-bot

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

@@ -238,6 +238,28 @@ def test_create_eval_run_with_inference_configs(client):
     assert evaluation_run.error is None
 
 
+def test_create_eval_run_with_metric_resource_name(client):
+    """Tests create_evaluation_run with metric_resource_name."""
+    client._api_client._http_options.api_version = "v1beta1"
+    client._api_client._http_options.base_url = (
+        "https://us-central1-staging-aiplatform.sandbox.googleapis.com/"
+    )
+    metric_resource_name = "projects/977012026409/locations/us-central1/evaluationMetrics/6048334299558576128"
+    metric = types.EvaluationRunMetric(
+        metric="my_custom_metric",
+        metric_resource_name=metric_resource_name,
+    )
+    evaluation_run = client.evals.create_evaluation_run(
+        dataset=types.EvaluationDataset(
+            eval_dataset_df=INPUT_DF_WITH_CONTEXT_AND_HISTORY
+        ),
+        metrics=[metric],
+        dest=GCS_DEST,
+    )
+    assert isinstance(evaluation_run, types.EvaluationRun)
+    assert evaluation_run.evaluation_config.metrics[0].metric == "my_custom_metric"
+
+
 # Dataframe tests fail in replay mode because of UUID generation mismatch.
 # def test_create_eval_run_data_source_evaluation_dataset(client):
 #     """Tests that create_evaluation_run() creates a correctly structured

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

@@ -353,6 +353,32 @@ def test_evaluation_agent_data(client):
         assert case_result.response_candidate_results is not None
 
 
+def test_metric_resource_name(client):
+    """Tests with a metric resource name in types.Metric."""
+    client._api_client._http_options.api_version = "v1beta1"
+    client._api_client._http_options.base_url = (
+        "https://us-central1-staging-aiplatform.sandbox.googleapis.com/"
+    )
+    metric_resource_name = "projects/977012026409/locations/us-central1/evaluationMetrics/6048334299558576128"
+    byor_df = pd.DataFrame(
+        {
+            "prompt": ["Write a simple story about a dinosaur"],
+            "response": ["Once upon a time, there was a T-Rex named Rexy."],
+        }
+    )
+    metric = types.Metric(
+        name="my_custom_metric", metric_resource_name=metric_resource_name
+    )
+    evaluation_result = client.evals.evaluate(
+        dataset=byor_df,
+        metrics=[metric],
+    )
+    assert isinstance(evaluation_result, types.EvaluationResult)
+    assert evaluation_result.eval_case_results is not None
+    assert len(evaluation_result.eval_case_results) > 0
+    assert evaluation_result.summary_metrics[0].metric_name == "my_custom_metric"
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),

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

@@ -143,19 +143,21 @@
 User prompt:
 {prompt}"""
 
+_PROMPTS_DF = pd.DataFrame(
+    {
+        "prompt": [
+            "Explain the theory of relativity in one sentence.",
+            "Write a short poem about a cat.",
+        ]
+    }
+)
+
 
 def test_public_method_generate_rubrics(client):
     """Tests the public generate_rubrics method."""
-    prompts_df = pd.DataFrame(
-        {
-            "prompt": [
-                "Explain the theory of relativity in one sentence.",
-                "Write a short poem about a cat.",
-            ]
-        }
-    )
+
     eval_dataset = client.evals.generate_rubrics(
-        src=prompts_df,
+        src=_PROMPTS_DF,
         prompt_template=_TEST_RUBRIC_GENERATION_PROMPT,
         rubric_group_name="text_quality_rubrics",
     )
@@ -176,6 +178,36 @@ def test_public_method_generate_rubrics(client):
     assert isinstance(first_rubric_group["text_quality_rubrics"][0], types.evals.Rubric)
 
 
+def test_public_method_generate_rubrics_with_metric(client):
+    """Tests the public generate_rubrics method with a metric."""
+    client._api_client._http_options.api_version = "v1beta1"
+    client._api_client._http_options.base_url = (
+        "https://us-central1-staging-aiplatform.sandbox.googleapis.com/"
+    )
+    metric_resource_name = "projects/977012026409/locations/us-central1/evaluationMetrics/6048334299558576128"
+    metric = types.Metric(
+        name="my_custom_metric", metric_resource_name=metric_resource_name
+    )
+    eval_dataset = client.evals.generate_rubrics(
+        src=_PROMPTS_DF, rubric_group_name="my_registered_rubrics", metric=metric
+    )
+    eval_dataset_df = eval_dataset.eval_dataset_df
+
+    assert isinstance(eval_dataset, types.EvaluationDataset)
+    assert isinstance(eval_dataset_df, pd.DataFrame)
+    assert "rubric_groups" in eval_dataset_df.columns
+    assert len(eval_dataset_df) == 2
+
+    first_rubric_group = eval_dataset_df["rubric_groups"][0]
+    assert isinstance(first_rubric_group, dict)
+    assert "my_registered_rubrics" in first_rubric_group
+    assert isinstance(first_rubric_group["my_registered_rubrics"], list)
+    assert first_rubric_group["my_registered_rubrics"]
+    assert isinstance(
+        first_rubric_group["my_registered_rubrics"][0], types.evals.Rubric
+    )
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),

vertexai/_genai/_evals_common.py

@@ -45,6 +45,7 @@
 from . import _gcs_utils
 from . import evals
 from . import types
+from . import _transformers as t
 
 logger = logging.getLogger(__name__)
 
@@ -1328,7 +1329,7 @@ def _resolve_dataset_inputs(
 
 
 def _resolve_evaluation_run_metrics(
-    metrics: list[types.EvaluationRunMetric], api_client: Any
+    metrics: Union[list[types.EvaluationRunMetric], list[types.Metric]], api_client: Any
 ) -> list[types.EvaluationRunMetric]:
     """Resolves a list of evaluation run metric instances, loading RubricMetric if necessary."""
     if not metrics:
@@ -1361,6 +1362,16 @@ def _resolve_evaluation_run_metrics(
                     e,
                 )
                 raise
+        elif isinstance(metric_instance, types.Metric):
+            config_dict = t.t_metrics([metric_instance])[0]
+            res_name = getattr(metric_instance, "metric_resource_name", None)
+            resolved_metrics_list.append(
+                types.EvaluationRunMetric(
+                    metric=metric_instance.name,
+                    metric_config=config_dict if config_dict else None,
+                    metric_resource_name=res_name,
+                )
+            )
         else:
             try:
                 metric_name_str = str(metric_instance)

vertexai/_genai/_evals_metric_handlers.py

@@ -1242,6 +1242,62 @@ def aggregate(
         )
 
 
+class RegisteredMetricHandler(PredefinedMetricHandler):
+    """Metric handler for registered metrics."""
+
+    def __init__(self, module: "evals.Evals", metric: types.Metric):
+        MetricHandler.__init__(self, module=module, metric=metric)
+
+    @override
+    def get_metric_result(
+        self, eval_case: types.EvalCase, response_index: int
+    ) -> types.EvalCaseMetricResult:
+        """Processes a single evaluation case for a registered metric."""
+        metric_name = self.metric.name
+        try:
+            payload = self._build_request_payload(eval_case, response_index)
+            for attempt in range(_MAX_RETRIES):
+                try:
+                    api_response = self.module._evaluate_instances(
+                        metric_sources=[self.metric],
+                        instance=payload.get("instance"),
+                        autorater_config=payload.get("autorater_config"),
+                    )
+                    break
+                except genai_errors.ClientError as e:
+                    if e.code == 429:
+                        if attempt == _MAX_RETRIES - 1:
+                            return types.EvalCaseMetricResult(
+                                metric_name=metric_name,
+                                error_message=f"Judge model resource exhausted after {_MAX_RETRIES} retries: {e}",
+                            )
+                        time.sleep(2**attempt)
+                    else:
+                        raise e
+
+            if api_response and api_response.metric_results:
+                result_data = api_response.metric_results[0]
+                error_message = None
+                if result_data.error and getattr(result_data.error, "code"):
+                    error_message = f"Error in metric result: {result_data.error}"
+                return types.EvalCaseMetricResult(
+                    metric_name=metric_name,
+                    score=result_data.score,
+                    explanation=result_data.explanation,
+                    rubric_verdicts=result_data.rubric_verdicts,
+                    error_message=error_message,
+                )
+            else:
+                return types.EvalCaseMetricResult(
+                    metric_name=metric_name,
+                    error_message="Metric results missing in API response.",
+                )
+        except Exception as e:
+            return types.EvalCaseMetricResult(
+                metric_name=metric_name, error_message=str(e)
+            )
+
+
 _METRIC_HANDLER_MAPPING = [
     (
         lambda m: hasattr(m, "remote_custom_function") and m.remote_custom_function,
@@ -1251,6 +1307,10 @@ def aggregate(
         lambda m: m.custom_function and isinstance(m.custom_function, Callable),
         CustomMetricHandler,
     ),
+    (
+        lambda m: getattr(m, "metric_resource_name", None) is not None,
+        RegisteredMetricHandler,
+    ),
     (
         lambda m: m.name in ComputationMetricHandler.SUPPORTED_COMPUTATION_METRICS,
         ComputationMetricHandler,

vertexai/_genai/_transformers.py

@@ -14,13 +14,16 @@
 #
 
 """Transformers module for Vertex addons."""
+import re
 from typing import Any
 
 from google.genai._common import get_value_by_path as getv
 
 from . import _evals_constant
 from . import types
 
+_METRIC_RES_NAME_RE = r"^projects/[^/]+/locations/[^/]+/evaluationMetrics/[^/]+$"
+
 
 def t_metrics(
     metrics: list["types.MetricSubclass"],
@@ -39,7 +42,8 @@ def t_metrics(
     for metric in metrics:
         metric_payload_item: dict[str, Any] = {}
 
-        metric_name = getv(metric, ["name"]).lower()
+        metric_id = getv(metric, ["metric"]) or getv(metric, ["name"])
+        metric_name = metric_id.lower() if metric_id else None
 
         if set_default_aggregation_metrics:
             metric_payload_item["aggregation_metrics"] = [
@@ -51,11 +55,13 @@ def t_metrics(
             metric_payload_item["exact_match_spec"] = {}
         elif metric_name == "bleu":
             metric_payload_item["bleu_spec"] = {}
-        elif metric_name.startswith("rouge"):
+        elif metric_name and metric_name.startswith("rouge"):
             rouge_type = metric_name.replace("_", "")
             metric_payload_item["rouge_spec"] = {"rouge_type": rouge_type}
         # API Pre-defined metrics
-        elif metric_name in _evals_constant.SUPPORTED_PREDEFINED_METRICS:
+        elif (
+            metric_name and metric_name in _evals_constant.SUPPORTED_PREDEFINED_METRICS
+        ):
             metric_payload_item["predefined_metric_spec"] = {
                 "metric_spec_name": metric_name,
                 "metric_spec_parameters": metric.metric_spec_parameters,
@@ -79,9 +85,38 @@ def t_metrics(
                     "return_raw_output": return_raw_output
                 }
             metric_payload_item["pointwise_metric_spec"] = pointwise_spec
+        elif getattr(metric, "metric_resource_name", None) is not None:
+            # Safe pass
+            pass
         else:
             raise ValueError(
                 f"Unsupported metric type or invalid metric name: {metric_name}"
             )
         metrics_payload.append(metric_payload_item)
     return metrics_payload
+
+
+def t_metric_sources(metrics: list[Any]) -> list[dict[str, Any]]:
+    """Prepares the MetricSource payload."""
+    sources_payload = []
+    for metric in metrics:
+        resource_name = getattr(metric, "metric_resource_name", None)
+        if (
+            not resource_name
+            and isinstance(metric, str)
+            and re.match(_METRIC_RES_NAME_RE, metric)
+        ):
+            resource_name = metric
+
+        if resource_name:
+            sources_payload.append({"metric_resource_name": resource_name})
+        else:
+            if hasattr(metric, "metric") and not isinstance(metric, str):
+                metric = metric.metric
+
+            if not hasattr(metric, "name"):
+                metric = types.Metric(name=str(metric))
+
+            metric_payload = t_metrics([metric])[0]
+            sources_payload.append({"metric": metric_payload})
+    return sources_payload