feat: Support per-eval case and per-invocation rubrics in rubric-based evaluators

Co-authored-by: Joseph Pagadora <[email protected]>
PiperOrigin-RevId: 853820099

Author: jcpagadora737

src/google/adk/cli/cli_eval.py

@@ -201,9 +201,11 @@ def pretty_print_eval_result(eval_result: EvalCaseResult):
           for r in metric_result.criterion.rubrics
       }
       for rubric_score in metric_result.details.rubric_scores:
-        rubric = rubrics_by_id.get(rubric_score.rubric_id)
+        rubric_text = rubrics_by_id.get(rubric_score.rubric_id)
+        if not rubric_text:
+          rubric_text = rubric_score.rubric_id
         click.echo(
-            f"Rubric: {rubric}, "
+            f"Rubric: {rubric_text}, "
             f"Score: {rubric_score.score}, "
             f"Reasoning: {rubric_score.rationale}"
         )
@@ -243,6 +245,8 @@ def pretty_print_eval_result(eval_result: EvalCaseResult):
         }
         for rubric_score in metric_result.details.rubric_scores:
           rubric = rubrics_by_id.get(rubric_score.rubric_id)
+          if not rubric:
+            rubric = rubric_score.rubric_id
           row_data[f"Rubric: {rubric}"] = (
               f"Reasoning: {rubric_score.rationale}, "
               f"Score: {rubric_score.score}"

src/google/adk/evaluation/local_eval_service.py

@@ -46,6 +46,7 @@
 from .eval_metrics import EvalMetricResult
 from .eval_metrics import EvalMetricResultDetails
 from .eval_metrics import EvalMetricResultPerInvocation
+from .eval_metrics import Rubric
 from .eval_result import EvalCaseResult
 from .eval_set import EvalCase
 from .eval_set_results_manager import EvalSetResultsManager
@@ -67,6 +68,46 @@ def _get_session_id() -> str:
   return f'{EVAL_SESSION_ID_PREFIX}{str(uuid.uuid4())}'
 
 
+def _add_rubrics_to_invocation(
+    invocation: Invocation, rubrics_to_add: list[Rubric]
+):
+  """Adds rubrics to invocation, throwing ValueError on duplicate rubric_id."""
+  if not invocation.rubrics:
+    invocation.rubrics = []
+  existing_ids = {r.rubric_id for r in invocation.rubrics}
+  for rubric in rubrics_to_add:
+    if rubric.rubric_id in existing_ids:
+      raise ValueError(
+          f"Rubric with rubric_id '{rubric.rubric_id}' already exists."
+      )
+    invocation.rubrics.append(rubric)
+    existing_ids.add(rubric.rubric_id)
+
+
+def _copy_eval_case_rubrics_to_actual_invocations(
+    eval_case: EvalCase, actual_invocations: list[Invocation]
+):
+  """Copies EvalCase level rubrics to all actual invocations."""
+  if hasattr(eval_case, 'rubrics') and eval_case.rubrics:
+    for invocation in actual_invocations:
+      _add_rubrics_to_invocation(invocation, eval_case.rubrics)
+
+
+def _copy_invocation_rubrics_to_actual_invocations(
+    expected_invocations: Optional[list[Invocation]],
+    actual_invocations: list[Invocation],
+):
+  """Copies invocation level rubrics to corresponding actual invocations."""
+  if expected_invocations:
+    for actual_invocation, expected_invocation in zip(
+        actual_invocations, expected_invocations
+    ):
+      if expected_invocation.rubrics:
+        _add_rubrics_to_invocation(
+            actual_invocation, expected_invocation.rubrics
+        )
+
+
 @experimental
 class LocalEvalService(BaseEvalService):
   """An implementation of BaseEvalService, that runs the evals locally."""
@@ -249,76 +290,27 @@ async def _evaluate_single_inference_result(
           )
       )
 
-    for eval_metric in evaluate_config.eval_metrics:
-      # Perform evaluation of the metric.
-      try:
-        with client_label_context(EVAL_CLIENT_LABEL):
-          evaluation_result = await self._evaluate_metric(
-              eval_metric=eval_metric,
-              actual_invocations=inference_result.inferences,
-              expected_invocations=eval_case.conversation,
-              conversation_scenario=eval_case.conversation_scenario,
-          )
-      except Exception as e:
-        # We intentionally catch the Exception as we don't want failures to
-        # affect other metric evaluation.
-        logger.error(
-            "Metric evaluation failed for metric `%s` for eval case id '%s'"
-            ' with following error `%s`',
-            eval_metric.metric_name,
-            eval_case.eval_id,
-            e,
-            exc_info=True,
-        )
-        # We use an empty result.
-        evaluation_result = EvaluationResult(
-            overall_eval_status=EvalStatus.NOT_EVALUATED
-        )
+    actual_invocations = inference_result.inferences
+    expected_invocations = eval_case.conversation
 
-      # Track overall score across all invocations.
-      eval_metric_result_details = EvalMetricResultDetails(
-          rubric_scores=evaluation_result.overall_rubric_scores
-      )
-      overall_eval_metric_results.append(
-          EvalMetricResult(
-              score=evaluation_result.overall_score,
-              eval_status=evaluation_result.overall_eval_status,
-              details=eval_metric_result_details,
-              **eval_metric.model_dump(),
-          )
-      )
+    # 1. Copy EvalCase level rubrics to all actual invocations.
+    _copy_eval_case_rubrics_to_actual_invocations(eval_case, actual_invocations)
 
-      if (
-          evaluation_result.overall_eval_status != EvalStatus.NOT_EVALUATED
-          and len(evaluation_result.per_invocation_results)
-          != len(eval_metric_result_per_invocation)
-      ):
-        raise ValueError(
-            'Eval metric should return results for each invocation. Found '
-            f'{len(evaluation_result.per_invocation_results)} results for '
-            f'{len(eval_metric_result_per_invocation)} invocations.'
-        )
+    # 2. If expected invocations are present, copy invocation level
+    # rubrics to corresponding actual invocations.
+    _copy_invocation_rubrics_to_actual_invocations(
+        expected_invocations, actual_invocations
+    )
 
-      # Track score across individual invocations.
-      for idx, invocation in enumerate(eval_metric_result_per_invocation):
-        invocation_result = (
-            evaluation_result.per_invocation_results[idx]
-            if evaluation_result.overall_eval_status != EvalStatus.NOT_EVALUATED
-            else PerInvocationResult(
-                actual_invocation=invocation.actual_invocation
-            )
-        )
-        eval_metric_result_details = EvalMetricResultDetails(
-            rubric_scores=invocation_result.rubric_scores
-        )
-        invocation.eval_metric_results.append(
-            EvalMetricResult(
-                score=invocation_result.score,
-                eval_status=invocation_result.eval_status,
-                details=eval_metric_result_details,
-                **eval_metric.model_dump(),
-            )
-        )
+    for eval_metric in evaluate_config.eval_metrics:
+      # Perform evaluation of the metric.
+      await self._evaluate_metric_for_eval_case(
+          eval_metric,
+          eval_case,
+          inference_result,
+          eval_metric_result_per_invocation,
+          overall_eval_metric_results,
+      )
 
     final_eval_status = self._generate_final_eval_status(
         overall_eval_metric_results
@@ -342,6 +334,84 @@ async def _evaluate_single_inference_result(
 
     return (inference_result, eval_case_result)
 
+  async def _evaluate_metric_for_eval_case(
+      self,
+      eval_metric: EvalMetric,
+      eval_case: EvalCase,
+      inference_result: InferenceResult,
+      eval_metric_result_per_invocation: list[EvalMetricResultPerInvocation],
+      overall_eval_metric_results: list[EvalMetricResult],
+  ):
+    """Performs evaluation of a metric for a given eval case and inference result."""
+    try:
+      with client_label_context(EVAL_CLIENT_LABEL):
+        evaluation_result = await self._evaluate_metric(
+            eval_metric=eval_metric,
+            actual_invocations=inference_result.inferences,
+            expected_invocations=eval_case.conversation,
+            conversation_scenario=eval_case.conversation_scenario,
+        )
+    except Exception as e:
+      # We intentionally catch the Exception as we don't want failures to
+      # affect other metric evaluation.
+      logger.error(
+          "Metric evaluation failed for metric `%s` for eval case id '%s'"
+          ' with following error `%s`',
+          eval_metric.metric_name,
+          eval_case.eval_id,
+          e,
+          exc_info=True,
+      )
+      # We use an empty result.
+      evaluation_result = EvaluationResult(
+          overall_eval_status=EvalStatus.NOT_EVALUATED
+      )
+
+    # Track overall score across all invocations.
+    eval_metric_result_details = EvalMetricResultDetails(
+        rubric_scores=evaluation_result.overall_rubric_scores
+    )
+    overall_eval_metric_results.append(
+        EvalMetricResult(
+            score=evaluation_result.overall_score,
+            eval_status=evaluation_result.overall_eval_status,
+            details=eval_metric_result_details,
+            **eval_metric.model_dump(),
+        )
+    )
+
+    if (
+        evaluation_result.overall_eval_status != EvalStatus.NOT_EVALUATED
+        and len(evaluation_result.per_invocation_results)
+        != len(eval_metric_result_per_invocation)
+    ):
+      raise ValueError(
+          'Eval metric should return results for each invocation. Found '
+          f'{len(evaluation_result.per_invocation_results)} results for '
+          f'{len(eval_metric_result_per_invocation)} invocations.'
+      )
+
+    # Track score across individual invocations.
+    for idx, invocation in enumerate(eval_metric_result_per_invocation):
+      invocation_result = (
+          evaluation_result.per_invocation_results[idx]
+          if evaluation_result.overall_eval_status != EvalStatus.NOT_EVALUATED
+          else PerInvocationResult(
+              actual_invocation=invocation.actual_invocation
+          )
+      )
+      eval_metric_result_details = EvalMetricResultDetails(
+          rubric_scores=invocation_result.rubric_scores
+      )
+      invocation.eval_metric_results.append(
+          EvalMetricResult(
+              score=invocation_result.score,
+              eval_status=invocation_result.eval_status,
+              details=eval_metric_result_details,
+              **eval_metric.model_dump(),
+          )
+      )
+
   async def _evaluate_metric(
       self,
       eval_metric: EvalMetric,

src/google/adk/evaluation/rubric_based_evaluator.py

@@ -328,28 +328,67 @@ def __init__(
     assert self._criterion.rubrics, "Rubrics are required."
 
     self._rubrics: list[Rubric] = self._criterion.rubrics
+    self._effective_rubrics_list: Optional[list[Rubric]] = None
 
     self._normalized_rubric_to_id_map = {
         _normalize_text(r.rubric_content.text_property): r.rubric_id
         for r in self._rubrics
     }
 
+  def create_effective_rubrics_list(
+      self,
+      invocation_rubrics: Optional[list[Rubric]],
+  ) -> None:
+    rubrics_by_id = {}
+
+    def _add_rubrics(rubrics_to_add: list[Rubric], scope_name: str):
+      for r in rubrics_to_add:
+        if r.rubric_id in rubrics_by_id:
+          raise ValueError(
+              f"Rubric with rubric_id '{r.rubric_id}' already exists. Rubric"
+              f" defined in {scope_name} conflicts with an existing rubric."
+          )
+        rubrics_by_id[r.rubric_id] = r
+
+    _add_rubrics(self._rubrics, "criterion")
+
+    if invocation_rubrics:
+      _add_rubrics(invocation_rubrics, "invocation")
+
+    self._effective_rubrics_list = list(rubrics_by_id.values())
+
+  def get_effective_rubrics_list(self) -> list[Rubric]:
+    """Returns the effective rubrics list."""
+    if self._effective_rubrics_list is None:
+      raise ValueError(
+          "Effective rubrics list not initialized. Call"
+          " create_effective_rubrics_list() first."
+      )
+    return self._effective_rubrics_list
+
   @override
   def convert_auto_rater_response_to_score(
-      self, auto_rater_response: LlmResponse
+      self,
+      auto_rater_response: LlmResponse,
   ) -> AutoRaterScore:
     """Returns an AutoRaterScore generated from AutoRater's response."""
     response_text = get_text_from_content(auto_rater_response.content)
     rubric_responses = self._auto_rater_response_parser.parse(response_text)
     rubric_scores = []
 
+    normalized_rubric_to_rubric_map = {}
+    for r in self.get_effective_rubrics_list():
+      normalized_rubric_to_rubric_map[
+          _normalize_text(r.rubric_content.text_property)
+      ] = r
+
     for rubric_response in rubric_responses:
-      normalized_rubric = _normalize_text(rubric_response.property_text)
-      rubric_id = self._normalized_rubric_to_id_map.get(normalized_rubric, None)
-      if rubric_id:
+      normalized_rubric_text = _normalize_text(rubric_response.property_text)
+      rubric = normalized_rubric_to_rubric_map.get(normalized_rubric_text, None)
+      if rubric:
         rubric_scores.append(
             RubricScore(
-                rubric_id=rubric_id,
+                rubric_id=rubric.rubric_id,
                 rationale=rubric_response.rationale,
                 score=rubric_response.score,
             )

src/google/adk/evaluation/rubric_based_final_response_quality_v1.py

@@ -25,6 +25,7 @@
 from .eval_case import InvocationEvents
 from .eval_metrics import EvalMetric
 from .eval_metrics import RubricsBasedCriterion
+from .eval_rubrics import Rubric
 from .llm_as_judge_utils import get_text_from_content
 from .llm_as_judge_utils import get_tool_calls_and_responses_as_json_str
 from .llm_as_judge_utils import get_tool_declarations_as_json_str
@@ -264,15 +265,19 @@ def __init__(self, eval_metric: EvalMetric):
 
   @override
   def format_auto_rater_prompt(
-      self, actual_invocation: Invocation, _: Optional[Invocation]
+      self,
+      actual_invocation: Invocation,
+      _: Optional[Invocation],
   ) -> str:
     """Returns the autorater prompt."""
-
+    self.create_effective_rubrics_list(actual_invocation.rubrics)
     user_input = get_text_from_content(actual_invocation.user_content)
     final_response = get_text_from_content(actual_invocation.final_response)
-    rubrics = "\n*  ".join(
-        [r.rubric_content.text_property for r in self._rubrics]
-    )
+
+    rubrics_text = "\n".join([
+        f"*  {r.rubric_content.text_property}"
+        for r in self._effective_rubrics_list
+    ])
 
     developer_instructions = ""
     tool_declarations = "Agent has no tools."
@@ -299,7 +304,7 @@ def format_auto_rater_prompt(
         user_input=user_input,
         response_steps=response_steps,
         final_response=final_response,
-        rubrics=rubrics,
+        rubrics=rubrics_text,
     )
 
     return auto_rater_prompt