fix: register custom metrics from eval config in LocalEvalSampler

Merge #6424

Closes #6177

PiperOrigin-RevId: 951011003

Author: ItsMacto

src/google/adk/cli/cli_eval.py

@@ -36,9 +36,6 @@
 from ..evaluation.eval_case import get_all_tool_calls
 from ..evaluation.eval_case import IntermediateDataType
 from ..evaluation.eval_metrics import EvalMetric
-from ..evaluation.eval_metrics import Interval
-from ..evaluation.eval_metrics import MetricInfo
-from ..evaluation.eval_metrics import MetricValueInfo
 from ..evaluation.eval_result import EvalCaseResult
 from ..evaluation.eval_sets_manager import EvalSetsManager
 from ..utils.context_utils import Aclosing
@@ -77,19 +74,6 @@ def _get_agent_module(agent_module_file_path: str) -> ModuleType:
   return _import_from_path(module_name, file_path)
 
 
-def get_default_metric_info(
-    metric_name: str, description: str = ""
-) -> MetricInfo:
-  """Returns a default MetricInfo for a metric."""
-  return MetricInfo(
-      metric_name=metric_name,
-      description=description,
-      metric_value_info=MetricValueInfo(
-          interval=Interval(min_value=0.0, max_value=1.0)
-      ),
-  )
-
-
 def get_root_agent(agent_module_file_path: str) -> Agent:
   """Returns root agent given the agent module."""
   agent_module = _get_agent_module(agent_module_file_path)

src/google/adk/cli/cli_tools_click.py

@@ -1005,7 +1005,6 @@ def cli_eval(
 
     from ..evaluation.base_eval_service import InferenceConfig
     from ..evaluation.base_eval_service import InferenceRequest
-    from ..evaluation.custom_metric_evaluator import _CustomMetricEvaluator
     from ..evaluation.eval_config import get_eval_metrics_from_config
     from ..evaluation.eval_config import get_evaluation_criteria_or_default
     from ..evaluation.evaluator import EvalStatus
@@ -1014,11 +1013,10 @@ def cli_eval(
     from ..evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
     from ..evaluation.local_eval_sets_manager import load_eval_set_from_file
     from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
-    from ..evaluation.metric_evaluator_registry import DEFAULT_METRIC_EVALUATOR_REGISTRY
+    from ..evaluation.metric_evaluator_registry import register_custom_metrics_from_config
     from ..evaluation.simulation.user_simulator_provider import UserSimulatorProvider
     from .cli_eval import _collect_eval_results
     from .cli_eval import _collect_inferences
-    from .cli_eval import get_default_metric_info
     from .cli_eval import get_root_agent
     from .cli_eval import parse_and_get_evals_to_run
     from .cli_eval import pretty_print_eval_result
@@ -1113,23 +1111,7 @@ def cli_eval(
   )
 
   try:
-    metric_evaluator_registry = DEFAULT_METRIC_EVALUATOR_REGISTRY
-    if eval_config.custom_metrics:
-      for (
-          metric_name,
-          config,
-      ) in eval_config.custom_metrics.items():
-        if config.metric_info:
-          metric_info = config.metric_info.model_copy()
-          metric_info.metric_name = metric_name
-        else:
-          metric_info = get_default_metric_info(
-              metric_name=metric_name, description=config.description
-          )
-
-        metric_evaluator_registry.register_evaluator(
-            metric_info, _CustomMetricEvaluator
-        )
+    metric_evaluator_registry = register_custom_metrics_from_config(eval_config)
 
     eval_service = LocalEvalService(
         root_agent=root_agent,

src/google/adk/evaluation/metric_evaluator_registry.py

@@ -15,12 +15,16 @@
 from __future__ import annotations
 
 import logging
+from typing import Optional
 
 from ..errors.not_found_error import NotFoundError
 from ..utils.feature_decorator import experimental
 from .custom_metric_evaluator import _CustomMetricEvaluator
+from .eval_config import EvalConfig
 from .eval_metrics import EvalMetric
+from .eval_metrics import Interval
 from .eval_metrics import MetricInfo
+from .eval_metrics import MetricValueInfo
 from .eval_metrics import PrebuiltMetrics
 from .evaluator import Evaluator
 from .final_response_match_v2 import FinalResponseMatchV2Evaluator
@@ -175,3 +179,51 @@ def _get_default_metric_evaluator_registry() -> MetricEvaluatorRegistry:
 
 
 DEFAULT_METRIC_EVALUATOR_REGISTRY = _get_default_metric_evaluator_registry()
+
+
+def _get_default_metric_info(
+    metric_name: str, description: str = ""
+) -> MetricInfo:
+  """Returns a default MetricInfo for a metric."""
+  return MetricInfo(
+      metric_name=metric_name,
+      description=description,
+      metric_value_info=MetricValueInfo(
+          interval=Interval(min_value=0.0, max_value=1.0)
+      ),
+  )
+
+
+def register_custom_metrics_from_config(
+    eval_config: EvalConfig,
+    metric_evaluator_registry: Optional[MetricEvaluatorRegistry] = None,
+) -> MetricEvaluatorRegistry:
+  """Registers custom metrics declared in the given eval config.
+
+  Args:
+    eval_config: The eval config whose custom_metrics entries should be
+      registered. Entries without a metric_info get a default one with a
+      [0.0, 1.0] value interval.
+    metric_evaluator_registry: The registry to register the metrics in.
+      Defaults to DEFAULT_METRIC_EVALUATOR_REGISTRY.
+
+  Returns:
+    The registry the metrics were registered in.
+  """
+  if metric_evaluator_registry is None:
+    metric_evaluator_registry = DEFAULT_METRIC_EVALUATOR_REGISTRY
+  if not eval_config.custom_metrics:
+    return metric_evaluator_registry
+
+  for metric_name, config in eval_config.custom_metrics.items():
+    if config.metric_info:
+      metric_info = config.metric_info.model_copy()
+      metric_info.metric_name = metric_name
+    else:
+      metric_info = _get_default_metric_info(
+          metric_name=metric_name, description=config.description
+      )
+    metric_evaluator_registry.register_evaluator(
+        metric_info, _CustomMetricEvaluator
+    )
+  return metric_evaluator_registry

src/google/adk/optimization/local_eval_sampler.py

@@ -38,6 +38,7 @@
 from ..evaluation.eval_result import EvalCaseResult
 from ..evaluation.eval_sets_manager import EvalSetsManager
 from ..evaluation.local_eval_service import LocalEvalService
+from ..evaluation.metric_evaluator_registry import register_custom_metrics_from_config
 from ..evaluation.simulation.user_simulator_provider import UserSimulatorProvider
 from ..utils.context_utils import Aclosing
 from .data_types import UnstructuredSamplingResult
@@ -154,6 +155,9 @@ def __init__(
   ):
     self._config = config
     self._eval_sets_manager = eval_sets_manager
+    self._metric_evaluator_registry = register_custom_metrics_from_config(
+        self._config.eval_config
+    )
 
     self._train_eval_set = self._config.train_eval_set
     self._train_eval_case_ids = (
@@ -236,6 +240,7 @@ async def _evaluate_agent(
     eval_service = LocalEvalService(
         root_agent=agent,
         eval_sets_manager=self._eval_sets_manager,
+        metric_evaluator_registry=self._metric_evaluator_registry,
         user_simulator_provider=user_simulator_provider,
     )
 

tests/unittests/evaluation/test_metric_evaluator_registry.py

@@ -14,17 +14,23 @@
 
 from __future__ import annotations
 
+from google.adk.agents.common_configs import CodeConfig
 from google.adk.errors.not_found_error import NotFoundError
+from google.adk.evaluation.custom_metric_evaluator import _CustomMetricEvaluator
+from google.adk.evaluation.eval_config import CustomMetricConfig
+from google.adk.evaluation.eval_config import EvalConfig
 from google.adk.evaluation.eval_metrics import EvalMetric
 from google.adk.evaluation.eval_metrics import Interval
 from google.adk.evaluation.eval_metrics import MetricInfo
 from google.adk.evaluation.eval_metrics import MetricValueInfo
 from google.adk.evaluation.eval_metrics import PrebuiltMetrics
 from google.adk.evaluation.evaluator import Evaluator
+from google.adk.evaluation.metric_evaluator_registry import DEFAULT_METRIC_EVALUATOR_REGISTRY
 from google.adk.evaluation.metric_evaluator_registry import FinalResponseMatchV2EvaluatorMetricInfoProvider
 from google.adk.evaluation.metric_evaluator_registry import HallucinationsV1EvaluatorMetricInfoProvider
 from google.adk.evaluation.metric_evaluator_registry import MetricEvaluatorRegistry
 from google.adk.evaluation.metric_evaluator_registry import PerTurnUserSimulatorQualityV1MetricInfoProvider
+from google.adk.evaluation.metric_evaluator_registry import register_custom_metrics_from_config
 from google.adk.evaluation.metric_evaluator_registry import ResponseEvaluatorMetricInfoProvider
 from google.adk.evaluation.metric_evaluator_registry import RubricBasedFinalResponseQualityV1EvaluatorMetricInfoProvider
 from google.adk.evaluation.metric_evaluator_registry import RubricBasedMultiTurnTrajectoryMetricInfoProvider
@@ -120,6 +126,112 @@ def test_get_evaluator_not_found(self, registry):
       registry.get_evaluator(eval_metric)
 
 
+class TestRegisterCustomMetricsFromConfig:
+  """Test cases for register_custom_metrics_from_config."""
+
+  _CUSTOM_METRIC_NAME = "custom_metric_for_registry_test"
+
+  @pytest.fixture
+  def registry(self):
+    registry = MetricEvaluatorRegistry()
+    yield registry
+    # The registry dict is shared class-level state; remove what we added.
+    registry._registry.pop(self._CUSTOM_METRIC_NAME, None)
+
+  def _registered_metric_info(self, registry, metric_name):
+    return next(
+        metric_info
+        for metric_info in registry.get_registered_metrics()
+        if metric_info.metric_name == metric_name
+    )
+
+  def test_registers_custom_metric_with_provided_metric_info(self, registry):
+    metric_info = MetricInfo(
+        metric_name="name_to_be_overridden",
+        description="Custom metric description",
+        metric_value_info=MetricValueInfo(
+            interval=Interval(min_value=0.0, max_value=5.0)
+        ),
+    )
+    eval_config = EvalConfig(
+        custom_metrics={
+            self._CUSTOM_METRIC_NAME: CustomMetricConfig(
+                code_config=CodeConfig(name="math.sqrt"),
+                metric_info=metric_info,
+            )
+        }
+    )
+
+    result = register_custom_metrics_from_config(eval_config, registry)
+
+    assert result is registry
+    registered_info = self._registered_metric_info(
+        registry, self._CUSTOM_METRIC_NAME
+    )
+    assert registered_info.metric_value_info.interval.max_value == 5.0
+    assert all(
+        metric_info.metric_name != "name_to_be_overridden"
+        for metric_info in registry.get_registered_metrics()
+    )
+    evaluator = registry.get_evaluator(
+        EvalMetric(
+            metric_name=self._CUSTOM_METRIC_NAME,
+            threshold=0.5,
+            custom_function_path="math.sqrt",
+        )
+    )
+    assert isinstance(evaluator, _CustomMetricEvaluator)
+
+  def test_registers_custom_metric_with_default_metric_info(self, registry):
+    eval_config = EvalConfig(
+        custom_metrics={
+            self._CUSTOM_METRIC_NAME: CustomMetricConfig(
+                code_config=CodeConfig(name="math.sqrt"),
+                description="A custom metric",
+            )
+        }
+    )
+
+    register_custom_metrics_from_config(eval_config, registry)
+
+    registered_info = self._registered_metric_info(
+        registry, self._CUSTOM_METRIC_NAME
+    )
+    assert registered_info.description == "A custom metric"
+    assert registered_info.metric_value_info.interval.min_value == 0.0
+    assert registered_info.metric_value_info.interval.max_value == 1.0
+
+  def test_no_custom_metrics_is_a_no_op(self, registry):
+    registered_before = registry.get_registered_metrics()
+
+    result = register_custom_metrics_from_config(EvalConfig(), registry)
+
+    assert result is registry
+    assert registry.get_registered_metrics() == registered_before
+
+  def test_defaults_to_the_default_registry(self):
+    eval_config = EvalConfig(
+        custom_metrics={
+            self._CUSTOM_METRIC_NAME: CustomMetricConfig(
+                code_config=CodeConfig(name="math.sqrt"),
+            )
+        }
+    )
+
+    try:
+      result = register_custom_metrics_from_config(eval_config)
+
+      assert result is DEFAULT_METRIC_EVALUATOR_REGISTRY
+      registered_info = self._registered_metric_info(
+          DEFAULT_METRIC_EVALUATOR_REGISTRY, self._CUSTOM_METRIC_NAME
+      )
+      assert registered_info.metric_name == self._CUSTOM_METRIC_NAME
+    finally:
+      DEFAULT_METRIC_EVALUATOR_REGISTRY._registry.pop(
+          self._CUSTOM_METRIC_NAME, None
+      )
+
+
 class TestMetricInfoProviders:
   """Test cases for MetricInfoProviders."""
 

tests/unittests/optimization/local_eval_sampler_test.py

@@ -14,22 +14,26 @@
 
 from __future__ import annotations
 
+from google.adk.agents.common_configs import CodeConfig
 from google.adk.agents.llm_agent import Agent
 from google.adk.evaluation.base_eval_service import EvaluateConfig
 from google.adk.evaluation.base_eval_service import EvaluateRequest
 from google.adk.evaluation.base_eval_service import InferenceConfig
 from google.adk.evaluation.base_eval_service import InferenceRequest
 from google.adk.evaluation.base_eval_service import InferenceResult
+from google.adk.evaluation.custom_metric_evaluator import _CustomMetricEvaluator
 from google.adk.evaluation.eval_case import Invocation
 from google.adk.evaluation.eval_case import InvocationEvent
 from google.adk.evaluation.eval_case import InvocationEvents
+from google.adk.evaluation.eval_config import CustomMetricConfig
 from google.adk.evaluation.eval_config import EvalConfig
 from google.adk.evaluation.eval_config import EvalMetric
 from google.adk.evaluation.eval_metrics import EvalMetricResult
 from google.adk.evaluation.eval_metrics import EvalMetricResultPerInvocation
 from google.adk.evaluation.eval_metrics import EvalStatus
 from google.adk.evaluation.eval_result import EvalCaseResult
 from google.adk.evaluation.eval_sets_manager import EvalSetsManager
+from google.adk.evaluation.metric_evaluator_registry import DEFAULT_METRIC_EVALUATOR_REGISTRY
 from google.adk.optimization.local_eval_sampler import _log_eval_summary
 from google.adk.optimization.local_eval_sampler import extract_single_invocation_info
 from google.adk.optimization.local_eval_sampler import extract_tool_call_data
@@ -218,6 +222,42 @@ def mock_get_eval_case_ids(self, eval_set_id):
     assert getattr(interface, attr) == expected_value
 
 
+def test_init_registers_custom_metrics(mocker):
+  mocker.patch.object(
+      LocalEvalSampler,
+      "_get_eval_case_ids",
+      autospec=True,
+      return_value=["t1"],
+  )
+  custom_metric_name = "custom_metric_for_sampler_test"
+  config = LocalEvalSamplerConfig(
+      eval_config=EvalConfig(
+          custom_metrics={
+              custom_metric_name: CustomMetricConfig(
+                  code_config=CodeConfig(name="math.sqrt")
+              )
+          }
+      ),
+      app_name="test_app",
+      train_eval_set="train_set",
+  )
+
+  try:
+    LocalEvalSampler(config, mocker.MagicMock(spec=EvalSetsManager))
+
+    evaluator = DEFAULT_METRIC_EVALUATOR_REGISTRY.get_evaluator(
+        EvalMetric(
+            metric_name=custom_metric_name,
+            threshold=0.5,
+            custom_function_path="math.sqrt",
+        )
+    )
+    assert isinstance(evaluator, _CustomMetricEvaluator)
+  finally:
+    # The registry dict is shared class-level state; remove what we added.
+    DEFAULT_METRIC_EVALUATOR_REGISTRY._registry.pop(custom_metric_name, None)
+
+
 @pytest.mark.asyncio
 async def test_evaluate_agent(mocker):
   # Mocking LocalEvalService and its methods