fix: scope custom metrics per registry and trust only the config path

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 956688159

Author: GWeale

src/google/adk/errors/not_found_error.py

@@ -18,7 +18,9 @@
 class NotFoundError(Exception):
   """Represents an error that occurs when an entity is not found."""
 
-  def __init__(self, message="The requested item was not found."):
+  def __init__(
+      self, message: str = "The requested item was not found."
+  ) -> None:
     """Initializes the NotFoundError exception.
 
     Args:

src/google/adk/evaluation/eval_config.py

@@ -269,27 +269,31 @@ def get_eval_metrics_from_config(eval_config: EvalConfig) -> list[EvalMetric]:
         custom_function_path = config.code_config.name
 
       if isinstance(criterion, float):
-        eval_metric_list.append(
-            EvalMetric(
-                metric_name=metric_name,
-                threshold=criterion,
-                criterion=BaseCriterion(threshold=criterion),
-                custom_function_path=custom_function_path,
-            )
+        eval_metric = EvalMetric(
+            metric_name=metric_name,
+            threshold=criterion,
+            criterion=BaseCriterion(threshold=criterion),
+            custom_function_path=custom_function_path,
         )
       elif isinstance(criterion, BaseCriterion):
-        eval_metric_list.append(
-            EvalMetric(
-                metric_name=metric_name,
-                threshold=criterion.threshold,
-                criterion=criterion,
-                custom_function_path=custom_function_path,
-            )
+        eval_metric = EvalMetric(
+            metric_name=metric_name,
+            threshold=criterion.threshold,
+            criterion=criterion,
+            custom_function_path=custom_function_path,
         )
       else:
         raise ValueError(
             f"Unexpected criterion type. {type(criterion).__name__} not"
             " supported."
         )
 
+      # The config is written by the developer running the eval, so the path it
+      # declares is the one honoured when the metric runs. It travels with the
+      # metric rather than in a registry keyed by metric name, so two apps in
+      # one process can declare the same metric name and each still gets its
+      # own function.
+      eval_metric._config_custom_function_path = custom_function_path  # pylint: disable=protected-access
+      eval_metric_list.append(eval_metric)
+
   return eval_metric_list

src/google/adk/evaluation/eval_metrics.py

@@ -25,6 +25,7 @@
 from pydantic import ConfigDict
 from pydantic import Field
 from pydantic import field_validator
+from pydantic import PrivateAttr
 from pydantic import SerializeAsAny
 from pydantic.json_schema import SkipJsonSchema
 from typing_extensions import TypeAlias
@@ -300,6 +301,11 @@ class EvalMetric(EvalBaseModel):
       description="""Path to custom function, if this is a custom metric.""",
   )
 
+  # The path declared for this metric in the eval config it was built from.
+  # Private, so that a metric parsed from an inbound payload cannot carry one:
+  # the public field above is settable by whoever built that payload.
+  _config_custom_function_path: Optional[str] = PrivateAttr(default=None)
+
 
 class EvalMetricResultDetails(EvalBaseModel):
   rubric_scores: Optional[list[RubricScore]] = Field(

src/google/adk/evaluation/metric_evaluator_registry.py

@@ -59,7 +59,15 @@
 class MetricEvaluatorRegistry:
   """A registry for metric Evaluators."""
 
-  _registry: dict[str, tuple[type[Evaluator], MetricInfo]] = {}
+  def __init__(self) -> None:
+    # Each registry instance owns its mappings, so a custom metric registered
+    # for one app is not resolvable from another app's registry. The standard
+    # metrics are seeded into every instance, as they are the same everywhere.
+    self._registry: dict[str, tuple[type[Evaluator], MetricInfo]] = {}
+    # Module path of the custom function backing a metric, keyed by metric
+    # name. Only ever written from an eval config.
+    self._custom_function_paths: dict[str, str] = {}
+    _register_standard_metrics(self)
 
   def get_evaluator(self, eval_metric: EvalMetric) -> Evaluator:
     """Returns an Evaluator for the given metric.
@@ -75,14 +83,34 @@ def get_evaluator(self, eval_metric: EvalMetric) -> Evaluator:
     if eval_metric.metric_name not in self._registry:
       raise NotFoundError(f"{eval_metric.metric_name} not found in registry.")
 
-    evaluator_type = self._registry[eval_metric.metric_name][0]
+    evaluator_type, _ = self._registry[eval_metric.metric_name]
     if issubclass(evaluator_type, _CustomMetricEvaluator):
+      custom_function_path = self._custom_function_path(eval_metric)
+      if custom_function_path is None:
+        raise NotFoundError(
+            f"No custom function registered for {eval_metric.metric_name}."
+        )
       return evaluator_type(
           eval_metric=eval_metric,
-          custom_function_path=eval_metric.custom_function_path,
+          custom_function_path=custom_function_path,
       )
     return evaluator_type(eval_metric=eval_metric)
 
+  def _custom_function_path(self, eval_metric: EvalMetric) -> Optional[str]:
+    """Returns the module path to import for a custom metric, if known.
+
+    Both sources are eval config entries: one recorded when the metric was
+    registered from a config, the other carried on a metric built from a
+    config. The `custom_function_path` field on the incoming metric is not
+    consulted, as it can be set by whoever built the request.
+
+    Args:
+      eval_metric: The metric whose custom function is being resolved.
+    """
+    if path := self._custom_function_paths.get(eval_metric.metric_name):
+      return path
+    return eval_metric._config_custom_function_path  # pylint: disable=protected-access
+
   def register_evaluator(
       self,
       metric_info: MetricInfo,
@@ -92,6 +120,25 @@ def register_evaluator(
 
     If a mapping already exist, then it is updated.
     """
+    self._register(metric_info, evaluator, custom_function_path=None)
+
+  def _register(
+      self,
+      metric_info: MetricInfo,
+      evaluator: type[Evaluator],
+      custom_function_path: Optional[str],
+  ) -> None:
+    """Registers an evaluator, along with the function path it may need.
+
+    A path already recorded for the metric is kept when this registration does
+    not carry one, so re-registering an evaluator does not drop it.
+
+    Args:
+      metric_info: Info for the metric the evaluator is registered against.
+      evaluator: The evaluator class to register.
+      custom_function_path: Module path of the function backing a custom
+        metric, taken from an eval config, or None.
+    """
     metric_name = metric_info.metric_name
     if metric_name in self._registry:
       logger.info(
@@ -102,6 +149,8 @@ def register_evaluator(
       )
 
     self._registry[str(metric_name)] = (evaluator, metric_info)
+    if custom_function_path is not None:
+      self._custom_function_paths[str(metric_name)] = custom_function_path
 
   def get_registered_metrics(
       self,
@@ -113,10 +162,10 @@ def get_registered_metrics(
     ]
 
 
-def _get_default_metric_evaluator_registry() -> MetricEvaluatorRegistry:
-  """Returns an instance of MetricEvaluatorRegistry with standard metrics already registered in it."""
-  metric_evaluator_registry = MetricEvaluatorRegistry()
-
+def _register_standard_metrics(
+    metric_evaluator_registry: MetricEvaluatorRegistry,
+) -> None:
+  """Registers the metrics that ship with ADK into the given registry."""
   metric_evaluator_registry.register_evaluator(
       metric_info=TrajectoryEvaluatorMetricInfoProvider().get_metric_info(),
       evaluator=TrajectoryEvaluator,
@@ -175,7 +224,10 @@ def _get_default_metric_evaluator_registry() -> MetricEvaluatorRegistry:
       evaluator=RubricBasedMultiTurnTrajectoryEvaluator,
   )
 
-  return metric_evaluator_registry
+
+def _get_default_metric_evaluator_registry() -> MetricEvaluatorRegistry:
+  """Returns an instance of MetricEvaluatorRegistry with standard metrics already registered in it."""
+  return MetricEvaluatorRegistry()
 
 
 DEFAULT_METRIC_EVALUATOR_REGISTRY = _get_default_metric_evaluator_registry()
@@ -223,7 +275,7 @@ def register_custom_metrics_from_config(
       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(  # pylint: disable=protected-access
+        metric_info, _CustomMetricEvaluator, config.code_config.name
     )
   return metric_evaluator_registry