chore: import upstream snapshot with attribution

tests/unittests/flows/llm_flows/test_tool_callbacks.py

@@ -0,0 +1,434 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+from google.adk.agents.llm_agent import Agent
+from google.adk.tools.base_tool import BaseTool
+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
+
+
+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
+
+  def __call__(
+      self,
+      tool: BaseTool,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+  ) -> dict[str, object]:
+    if self.modify_tool_request:
+      args['input_str'] = 'modified_input'
+      return None
+    return self.mock_response
+
+
+class MockAfterToolCallback(BaseModel):
+  """Mock after tool callback."""
+
+  mock_response: dict[str, object]
+  modify_tool_request: bool = False
+  modify_tool_response: bool = False
+
+  def __call__(
+      self,
+      tool: BaseTool,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      tool_response: dict[str, Any] = None,
+  ) -> dict[str, object]:
+    if self.modify_tool_request:
+      args['input_str'] = 'modified_input'
+      return None
+    if self.modify_tool_response:
+      tool_response['result'] = 'modified_output'
+      return tool_response
+    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',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_tool_callback=MockBeforeToolCallback(
+          mock_response={'test': 'before_tool_callback'}
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', Part.from_function_call(name='simple_function', args={})),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function', response={'test': 'before_tool_callback'}
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+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'}
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_tool_callback=noop_callback,
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function', args={'input_str': 'simple_function_call'}
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function',
+              response={'result': 'simple_function_call'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+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',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_tool_callback=MockBeforeToolCallback(
+          mock_response={'test': 'before_tool_callback'},
+          modify_tool_request=True,
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', Part.from_function_call(name='simple_function', args={})),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function',
+              response={'result': 'modified_input'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+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'}
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_tool_callback=MockAfterToolCallback(
+          mock_response={'test': 'after_tool_callback'}
+      ),
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function', args={'input_str': 'simple_function_call'}
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function', response={'test': 'after_tool_callback'}
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+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'}
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_tool_callback=noop_callback,
+      tools=[simple_function],
+  )
+
+  runner = testing_utils.InMemoryRunner(agent)
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      (
+          'root_agent',
+          Part.from_function_call(
+              name='simple_function', args={'input_str': 'simple_function_call'}
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function',
+              response={'result': 'simple_function_call'},
+          ),
+      ),
+      ('root_agent', 'response1'),
+  ]
+
+
+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'}
+      ),
+      'response1',
+  ]
+  mock_model = testing_utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_tool_callback=MockAfterToolCallback(
+          mock_response={'result': 'after_tool_callback'},
+          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='simple_function', args={'input_str': 'simple_function_call'}
+          ),
+      ),
+      (
+          'root_agent',
+          Part.from_function_response(
+              name='simple_function',
+              response={'result': 'modified_output'},
+          ),
+      ),
+      ('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'),
+  ]