diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 2fa3ef56..a02f97dd 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -105,6 +105,16 @@ AfterToolCallback: TypeAlias = Union[ 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 @@ class LlmAgent(BaseAgent): 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 ( diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 40e0be25..91f8808f 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/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 _execute_single_function_call_async( 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 diff --git a/tests/unittests/flows/llm_flows/test_tool_callbacks.py b/tests/unittests/flows/llm_flows/test_tool_callbacks.py index 59845b61..b839a3c9 100644 --- a/tests/unittests/flows/llm_flows/test_tool_callbacks.py +++ b/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'), + ]