feat: Add on_tool_error_callback in LlmAgent

PiperOrigin-RevId: 824598502

Author: xuanyang15

src/google/adk/agents/llm_agent.py

@@ -105,6 +105,16 @@
     list[_SingleAfterToolCallback],
 ]
 
+_SingleOnToolErrorCallback: TypeAlias = Callable[
+    [BaseTool, dict[str, Any], ToolContext, Exception],
+    Union[Awaitable[Optional[dict]], Optional[dict]],
+]
+
+OnToolErrorCallback: TypeAlias = Union[
+    _SingleOnToolErrorCallback,
+    list[_SingleOnToolErrorCallback],
+]
+
 InstructionProvider: TypeAlias = Callable[
     [ReadonlyContext], Union[str, Awaitable[str]]
 ]
@@ -381,6 +391,21 @@ class LlmAgent(BaseAgent):
     tool_context: ToolContext,
     tool_response: The response from the tool.
 
+  Returns:
+    When present, the returned dict will be used as tool result.
+  """
+  on_tool_error_callback: Optional[OnToolErrorCallback] = None
+  """Callback or list of callbacks to be called when a tool call encounters an error.
+
+  When a list of callbacks is provided, the callbacks will be called in the
+  order they are listed until a callback does not return None.
+
+  Args:
+    tool: The tool to be called.
+    args: The arguments to the tool.
+    tool_context: ToolContext,
+    error: The error from the tool call.
+
   Returns:
     When present, the returned dict will be used as tool result.
   """
@@ -582,6 +607,20 @@ def canonical_after_tool_callbacks(
       return self.after_tool_callback
     return [self.after_tool_callback]
 
+  @property
+  def canonical_on_tool_error_callbacks(
+      self,
+  ) -> list[OnToolErrorCallback]:
+    """The resolved self.on_tool_error_callback field as a list of OnToolErrorCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.on_tool_error_callback:
+      return []
+    if isinstance(self.on_tool_error_callback, list):
+      return self.on_tool_error_callback
+    return [self.on_tool_error_callback]
+
   @property
   def _llm_flow(self) -> BaseLlmFlow:
     if (

src/google/adk/flows/llm_flows/functions.py

@@ -275,6 +275,40 @@ async def _execute_single_function_call_async(
     tool_confirmation: Optional[ToolConfirmation] = None,
 ) -> Optional[Event]:
   """Execute a single function call with thread safety for state modifications."""
+
+  async def _run_on_tool_error_callbacks(
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+      error: Exception,
+  ) -> Optional[dict[str, Any]]:
+    """Runs the on_tool_error_callbacks for the given tool."""
+    error_response = (
+        await invocation_context.plugin_manager.run_on_tool_error_callback(
+            tool=tool,
+            tool_args=tool_args,
+            tool_context=tool_context,
+            error=error,
+        )
+    )
+    if error_response is not None:
+      return error_response
+
+    for callback in agent.canonical_on_tool_error_callbacks:
+      error_response = callback(
+          tool=tool,
+          args=tool_args,
+          tool_context=tool_context,
+          error=error,
+      )
+      if inspect.isawaitable(error_response):
+        error_response = await error_response
+      if error_response is not None:
+        return error_response
+
+    return None
+
   # Do not use "args" as the variable name, because it is a reserved keyword
   # in python debugger.
   # Make a deep copy to avoid being modified.
@@ -290,13 +324,11 @@ async def _execute_single_function_call_async(
     tool = _get_tool(function_call, tools_dict)
   except ValueError as tool_error:
     tool = BaseTool(name=function_call.name, description='Tool not found')
-    error_response = (
-        await invocation_context.plugin_manager.run_on_tool_error_callback(
-            tool=tool,
-            tool_args=function_args,
-            tool_context=tool_context,
-            error=tool_error,
-        )
+    error_response = await _run_on_tool_error_callbacks(
+        tool=tool,
+        tool_args=function_args,
+        tool_context=tool_context,
+        error=tool_error,
     )
     if error_response is not None:
       return __build_response_event(
@@ -335,13 +367,11 @@ async def _run_with_trace():
             tool, args=function_args, tool_context=tool_context
         )
       except Exception as tool_error:
-        error_response = (
-            await invocation_context.plugin_manager.run_on_tool_error_callback(
-                tool=tool,
-                tool_args=function_args,
-                tool_context=tool_context,
-                error=tool_error,
-            )
+        error_response = await _run_on_tool_error_callbacks(
+            tool=tool,
+            tool_args=function_args,
+            tool_context=tool_context,
+            error=tool_error,
         )
         if error_response is not None:
           function_response = error_response

tests/unittests/flows/llm_flows/test_tool_callbacks.py

@@ -20,6 +20,7 @@
 from google.genai import types
 from google.genai.types import Part
 from pydantic import BaseModel
+import pytest
 
 from ... import testing_utils
 
@@ -28,7 +29,13 @@ def simple_function(input_str: str) -> str:
   return {'result': input_str}
 
 
+def simple_function_with_error() -> str:
+  raise SystemError('simple_function_with_error')
+
+
 class MockBeforeToolCallback(BaseModel):
+  """Mock before tool callback."""
+
   mock_response: dict[str, object]
   modify_tool_request: bool = False
 
@@ -45,6 +52,8 @@ def __call__(
 
 
 class MockAfterToolCallback(BaseModel):
+  """Mock after tool callback."""
+
   mock_response: dict[str, object]
   modify_tool_request: bool = False
   modify_tool_response: bool = False
@@ -65,13 +74,32 @@ def __call__(
     return self.mock_response
 
 
+class MockOnToolErrorCallback(BaseModel):
+  """Mock on tool error callback."""
+
+  mock_response: dict[str, object]
+  modify_tool_response: bool = False
+
+  def __call__(
+      self,
+      tool: BaseTool,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      error: Exception,
+  ) -> dict[str, object]:
+    if self.modify_tool_response:
+      return self.mock_response
+    return None
+
+
 def noop_callback(
     **kwargs,
 ) -> dict[str, object]:
   pass
 
 
 def test_before_tool_callback():
+  """Test that the before_tool_callback is called before the tool is called."""
   responses = [
       types.Part.from_function_call(name='simple_function', args={}),
       'response1',
@@ -100,6 +128,7 @@ def test_before_tool_callback():
 
 
 def test_before_tool_callback_noop():
+  """Test that the before_tool_callback is a no-op when not overridden."""
   responses = [
       types.Part.from_function_call(
           name='simple_function', args={'input_str': 'simple_function_call'}
@@ -134,6 +163,7 @@ def test_before_tool_callback_noop():
 
 
 def test_before_tool_callback_modify_tool_request():
+  """Test that the before_tool_callback modifies the tool request."""
   responses = [
       types.Part.from_function_call(name='simple_function', args={}),
       'response1',
@@ -164,6 +194,7 @@ def test_before_tool_callback_modify_tool_request():
 
 
 def test_after_tool_callback():
+  """Test that the after_tool_callback is called after the tool is called."""
   responses = [
       types.Part.from_function_call(
           name='simple_function', args={'input_str': 'simple_function_call'}
@@ -199,6 +230,7 @@ def test_after_tool_callback():
 
 
 def test_after_tool_callback_noop():
+  """Test that the after_tool_callback is a no-op when not overridden."""
   responses = [
       types.Part.from_function_call(
           name='simple_function', args={'input_str': 'simple_function_call'}
@@ -233,6 +265,7 @@ def test_after_tool_callback_noop():
 
 
 def test_after_tool_callback_modify_tool_response():
+  """Test that the after_tool_callback modifies the tool response."""
   responses = [
       types.Part.from_function_call(
           name='simple_function', args={'input_str': 'simple_function_call'}
@@ -267,3 +300,135 @@ def test_after_tool_callback_modify_tool_response():
       ),
       ('root_agent', 'response1'),
   ]
+
+
+async def test_on_tool_error_callback_tool_not_found_noop():
+  """Test that the on_tool_error_callback is a no-op when the tool is not found."""
+  responses = [
+      types.Part.from_function_call(
+          name='nonexistent_function',
+          args={'input_str': 'simple_function_call'},
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_tool_error_callback=noop_callback,
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  with pytest.raises(ValueError):
+    await runner.run_async('test')
+
+
+def test_on_tool_error_callback_tool_not_found_modify_tool_response():
+  """Test that the on_tool_error_callback modifies the tool response when the tool is not found."""
+  responses = [
+      types.Part.from_function_call(
+          name='nonexistent_function',
+          args={'input_str': 'simple_function_call'},
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_tool_error_callback=MockOnToolErrorCallback(
+          mock_response={'result': 'on_tool_error_callback_response'},
+          modify_tool_response=True,
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='nonexistent_function',
+              args={'input_str': 'simple_function_call'},
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='nonexistent_function',
+              response={'result': 'on_tool_error_callback_response'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+async def test_on_tool_error_callback_tool_error_noop():
+  """Test that the on_tool_error_callback is a no-op when the tool returns an error."""
+  responses = [
+      types.Part.from_function_call(
+          name='simple_function_with_error',
+          args={},
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_tool_error_callback=noop_callback,
+      tools=[simple_function_with_error],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  with pytest.raises(SystemError):
+    await runner.run_async('test')
+
+
+def test_on_tool_error_callback_tool_error_modify_tool_response():
+  """Test that the on_tool_error_callback modifies the tool response when the tool returns an error."""
+
+  async def async_on_tool_error_callback(
+      tool: BaseTool,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      error: Exception,
+  ) -> dict[str, object]:
+    if tool.name == 'simple_function_with_error':
+      return {'result': 'async_on_tool_error_callback_response'}
+    return None
+
+  responses = [
+      types.Part.from_function_call(
+          name='simple_function_with_error',
+          args={},
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_tool_error_callback=async_on_tool_error_callback,
+      tools=[simple_function_with_error],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function_with_error',
+              args={},
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function_with_error',
+              response={'result': 'async_on_tool_error_callback_response'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]