feat: Add on_tool_error_callback in LlmAgent

PiperOrigin-RevId: 824598502

tests/unittests/flows/llm_flows/test_tool_callbacks.py

@@ -20,6 +20,7 @@ from google.adk.tools.tool_context import ToolContext
 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 @@ class MockBeforeToolCallback(BaseModel):
 
 
 class MockAfterToolCallback(BaseModel):
+  """Mock after tool callback."""
+
   mock_response: dict[str, object]
   modify_tool_request: bool = False
   modify_tool_response: bool = False
@@ -65,6 +74,24 @@ class MockAfterToolCallback(BaseModel):
     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]:
@@ -72,6 +99,7 @@ def noop_callback(
 
 
 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'),
+  ]