fix: Fix exception handling and argument order in ReflectRetryToolPlugin

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 906753501

Author: xuanyang15

src/google/adk/optimization/local_eval_sampler.py

@@ -289,11 +289,7 @@ def _extract_eval_data(
         for eval_metric_result in per_invocation_result.eval_metric_results:
           eval_metric_results.append({
               "metric_name": eval_metric_result.metric_name,
-              "score": (
-                  round(eval_metric_result.score, 2)
-                  if eval_metric_result.score is not None
-                  else None
-              ),  # accurate enough
+              "score": round(eval_metric_result.score, 2),  # accurate enough
               "eval_status": eval_metric_result.eval_status.name,
           })
         per_invocation_result_dict = {

src/google/adk/plugins/reflect_retry_tool_plugin.py

@@ -242,8 +242,8 @@ async def _handle_tool_error(
     """
     if self.max_retries == 0:
       if self.throw_exception_if_retry_exceeded:
-        raise error
-      return self._get_tool_retry_exceed_msg(tool, error, tool_args)
+        raise self._ensure_exception(error)
+      return self._get_tool_retry_exceed_msg(tool, tool_args, error)
 
     scope_key = self._get_scope_key(tool_context)
     async with self._lock:
@@ -260,7 +260,7 @@ async def _handle_tool_error(
 
       # Max Retry exceeded
       if self.throw_exception_if_retry_exceeded:
-        raise error
+        raise self._ensure_exception(error)
       else:
         return self._get_tool_retry_exceed_msg(tool, tool_args, error)
 

tests/unittests/optimization/local_eval_sampler_test.py

@@ -338,48 +338,6 @@ async def test_extract_eval_data(mocker):
   ]
 
 
-def test_extract_eval_data_preserves_none_metric_score(mocker):
-  mock_eval_sets_manager = mocker.MagicMock(spec=EvalSetsManager)
-  mock_eval_case = mocker.MagicMock()
-  mock_eval_case.conversation_scenario = "test_scenario"
-  mock_eval_sets_manager.get_eval_case.return_value = mock_eval_case
-
-  mock_metric_result = mocker.MagicMock(spec=EvalMetricResult)
-  mock_metric_result.metric_name = "test_metric"
-  mock_metric_result.score = None
-  mock_metric_result.eval_status = EvalStatus.NOT_EVALUATED
-
-  mock_per_inv_result = mocker.MagicMock(spec=EvalMetricResultPerInvocation)
-  mock_per_inv_result.actual_invocation = mocker.MagicMock(spec=Invocation)
-  mock_per_inv_result.expected_invocation = mocker.MagicMock(spec=Invocation)
-  mock_per_inv_result.eval_metric_results = [mock_metric_result]
-
-  mock_eval_result = mocker.MagicMock(spec=EvalCaseResult)
-  mock_eval_result.eval_id = "t1"
-  mock_eval_result.eval_metric_result_per_invocation = [mock_per_inv_result]
-
-  mocker.patch(
-      "google.adk.optimization.local_eval_sampler.extract_single_invocation_info",
-      side_effect=[{"info": "actual"}, {"info": "expected"}],
-  )
-
-  config = LocalEvalSamplerConfig(
-      eval_config=EvalConfig(),
-      app_name="test_app",
-      train_eval_set="train_set",
-      train_eval_case_ids=["t1"],
-  )
-  interface = LocalEvalSampler(config, mock_eval_sets_manager)
-
-  eval_data = interface._extract_eval_data("train_set", [mock_eval_result])
-
-  assert eval_data["t1"]["invocations"][0]["eval_metric_results"] == [{
-      "metric_name": "test_metric",
-      "score": None,
-      "eval_status": "NOT_EVALUATED",
-  }]
-
-
 @pytest.mark.asyncio
 async def test_sample_and_score(mocker):
   # Mock results

tests/unittests/plugins/test_reflect_retry_tool_plugin.py

@@ -168,6 +168,57 @@ async def test_on_tool_error_callback_max_retries_zero(self):
     # Should re-raise the original exception when max_retries is 0
     self.assertIs(cm.exception, error)
 
+  async def test_on_tool_error_callback_max_retries_zero_without_exception(
+      self,
+  ):
+    """Test error callback when max_retries is 0 and exception is disabled."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=0, throw_exception_if_retry_exceeded=False
+    )
+    error = ValueError("Test error")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Should return a retry exceeded message instead of raising
+    self.assertIsNotNone(result)
+    self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result["error_type"], "ValueError")
+    self.assertEqual(result["retry_count"], 0)
+    self.assertIn(
+        "the retry limit has been exceeded", result["reflection_guidance"]
+    )
+
+  async def test_on_tool_error_callback_max_retries_zero_with_dict_error(self):
+    """Test error callback when max_retries is 0 and error is a dict."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = CustomErrorExtractionPlugin(
+        max_retries=0, throw_exception_if_retry_exceeded=True
+    )
+    dict_error = {"status": "error", "message": "Custom dict error"}
+    plugin.set_error_condition(lambda result: dict_error)
+
+    with self.assertRaises(Exception) as cm:
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          result={"some": "result"},
+      )
+
+    # Should raise an Exception wrapping the dict
+    self.assertNotIsInstance(cm.exception, TypeError)
+    self.assertIn("Custom dict error", str(cm.exception))
+
   async def test_on_tool_error_callback_first_failure(self):
     """Test first tool failure creates reflection response."""
     plugin = self.get_plugin()
@@ -280,6 +331,40 @@ async def test_max_retries_exceeded_with_exception(self):
     # Verify exception properties
     self.assertIs(cm.exception, error)
 
+  async def test_max_retries_exceeded_with_dict_error(self):
+    """Test that Exception is raised when max retries exceeded with dict error."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = CustomErrorExtractionPlugin(
+        max_retries=1, throw_exception_if_retry_exceeded=True
+    )
+    dict_error = {"status": "error", "message": "Custom dict error"}
+    plugin.set_error_condition(lambda result: dict_error)
+
+    # First call should fail and return a retry response
+    result1 = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result={"some": "result"},
+    )
+    self.assertIsNotNone(result1)
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Second call should exceed max_retries and raise
+    with self.assertRaises(Exception) as cm:
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          result={"some": "result"},
+      )
+
+    # Verify exception properties
+    self.assertNotIsInstance(cm.exception, TypeError)
+    self.assertIn("Custom dict error", str(cm.exception))
+
   async def test_max_retries_exceeded_without_exception(self):
     """Test max retries exceeded returns failure message when exception is disabled."""
     mock_tool = self.get_mock_tool()