Support async agent and model callbacks

PiperOrigin-RevId: 755542756

tests/unittests/agents/__init__.py

@@ -11,4 +11,3 @@
 # 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.
-

tests/unittests/agents/test_base_agent.py

@@ -33,16 +33,34 @@ def _before_agent_callback_noop(callback_context: CallbackContext) -> None:
   pass
 
 
+async def _async_before_agent_callback_noop(
+    callback_context: CallbackContext,
+) -> None:
+  pass
+
+
 def _before_agent_callback_bypass_agent(
     callback_context: CallbackContext,
 ) -> types.Content:
   return types.Content(parts=[types.Part(text='agent run is bypassed.')])
 
 
+async def _async_before_agent_callback_bypass_agent(
+    callback_context: CallbackContext,
+) -> types.Content:
+  return types.Content(parts=[types.Part(text='agent run is bypassed.')])
+
+
 def _after_agent_callback_noop(callback_context: CallbackContext) -> None:
   pass
 
 
+async def _async_after_agent_callback_noop(
+    callback_context: CallbackContext,
+) -> None:
+  pass
+
+
 def _after_agent_callback_append_agent_reply(
     callback_context: CallbackContext,
 ) -> types.Content:
@@ -51,6 +69,14 @@ def _after_agent_callback_append_agent_reply(
   )
 
 
+async def _async_after_agent_callback_append_agent_reply(
+    callback_context: CallbackContext,
+) -> types.Content:
+  return types.Content(
+      parts=[types.Part(text='Agent reply from after agent callback.')]
+  )
+
+
 class _IncompleteAgent(BaseAgent):
   pass
 
@@ -158,6 +184,34 @@ async def test_run_async_before_agent_callback_noop(
   spy_run_async_impl.assert_called_once()
 
 
+@pytest.mark.asyncio
+async def test_run_async_with_async_before_agent_callback_noop(
+    request: pytest.FixtureRequest,
+    mocker: pytest_mock.MockerFixture,
+) -> Union[types.Content, None]:
+  # Arrange
+  agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent',
+      before_agent_callback=_async_before_agent_callback_noop,
+  )
+  parent_ctx = _create_parent_invocation_context(
+      request.function.__name__, agent
+  )
+  spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__)
+  spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback')
+
+  # Act
+  _ = [e async for e in agent.run_async(parent_ctx)]
+
+  # Assert
+  spy_before_agent_callback.assert_called_once()
+  _, kwargs = spy_before_agent_callback.call_args
+  assert 'callback_context' in kwargs
+  assert isinstance(kwargs['callback_context'], CallbackContext)
+
+  spy_run_async_impl.assert_called_once()
+
+
 @pytest.mark.asyncio
 async def test_run_async_before_agent_callback_bypass_agent(
     request: pytest.FixtureRequest,
@@ -185,6 +239,33 @@ async def test_run_async_before_agent_callback_bypass_agent(
   assert events[0].content.parts[0].text == 'agent run is bypassed.'
 
 
+@pytest.mark.asyncio
+async def test_run_async_with_async_before_agent_callback_bypass_agent(
+    request: pytest.FixtureRequest,
+    mocker: pytest_mock.MockerFixture,
+):
+  # Arrange
+  agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent',
+      before_agent_callback=_async_before_agent_callback_bypass_agent,
+  )
+  parent_ctx = _create_parent_invocation_context(
+      request.function.__name__, agent
+  )
+  spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__)
+  spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback')
+
+  # Act
+  events = [e async for e in agent.run_async(parent_ctx)]
+
+  # Assert
+  spy_before_agent_callback.assert_called_once()
+  spy_run_async_impl.assert_not_called()
+
+  assert len(events) == 1
+  assert events[0].content.parts[0].text == 'agent run is bypassed.'
+
+
 @pytest.mark.asyncio
 async def test_run_async_after_agent_callback_noop(
     request: pytest.FixtureRequest,
@@ -211,6 +292,32 @@ async def test_run_async_after_agent_callback_noop(
   assert len(events) == 1
 
 
+@pytest.mark.asyncio
+async def test_run_async_with_async_after_agent_callback_noop(
+    request: pytest.FixtureRequest,
+    mocker: pytest_mock.MockerFixture,
+):
+  # Arrange
+  agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent',
+      after_agent_callback=_async_after_agent_callback_noop,
+  )
+  parent_ctx = _create_parent_invocation_context(
+      request.function.__name__, agent
+  )
+  spy_after_agent_callback = mocker.spy(agent, 'after_agent_callback')
+
+  # Act
+  events = [e async for e in agent.run_async(parent_ctx)]
+
+  # Assert
+  spy_after_agent_callback.assert_called_once()
+  _, kwargs = spy_after_agent_callback.call_args
+  assert 'callback_context' in kwargs
+  assert isinstance(kwargs['callback_context'], CallbackContext)
+  assert len(events) == 1
+
+
 @pytest.mark.asyncio
 async def test_run_async_after_agent_callback_append_reply(
     request: pytest.FixtureRequest,
@@ -236,6 +343,31 @@ async def test_run_async_after_agent_callback_append_reply(
   )
 
 
+@pytest.mark.asyncio
+async def test_run_async_with_async_after_agent_callback_append_reply(
+    request: pytest.FixtureRequest,
+):
+  # Arrange
+  agent = _TestingAgent(
+      name=f'{request.function.__name__}_test_agent',
+      after_agent_callback=_async_after_agent_callback_append_agent_reply,
+  )
+  parent_ctx = _create_parent_invocation_context(
+      request.function.__name__, agent
+  )
+
+  # Act
+  events = [e async for e in agent.run_async(parent_ctx)]
+
+  # Assert
+  assert len(events) == 2
+  assert events[1].author == agent.name
+  assert (
+      events[1].content.parts[0].text
+      == 'Agent reply from after agent callback.'
+  )
+
+
 @pytest.mark.asyncio
 async def test_run_async_incomplete_agent(request: pytest.FixtureRequest):
   agent = _IncompleteAgent(name=f'{request.function.__name__}_test_agent')

tests/unittests/agents/test_llm_agent_callbacks.py

@@ -56,10 +56,44 @@ class MockAfterModelCallback(BaseModel):
     )
 
 
+class MockAsyncBeforeModelCallback(BaseModel):
+  mock_response: str
+
+  async def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
+class MockAsyncAfterModelCallback(BaseModel):
+  mock_response: str
+
+  async def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_response: LlmResponse,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
 def noop_callback(**kwargs) -> Optional[LlmResponse]:
   pass
 
 
+async def async_noop_callback(**kwargs) -> Optional[LlmResponse]:
+  pass
+
+
 @pytest.mark.asyncio
 async def test_before_model_callback():
   responses = ['model_response']
@@ -98,26 +132,6 @@ async def test_before_model_callback_noop():
   ]
 
 
-@pytest.mark.asyncio
-async def test_before_model_callback_end():
-  responses = ['model_response']
-  mock_model = utils.MockModel.create(responses=responses)
-  agent = Agent(
-      name='root_agent',
-      model=mock_model,
-      before_model_callback=MockBeforeModelCallback(
-          mock_response='before_model_callback',
-      ),
-  )
-
-  runner = utils.TestInMemoryRunner(agent)
-  assert utils.simplify_events(
-      await runner.run_async_with_new_session('test')
-  ) == [
-      ('root_agent', 'before_model_callback'),
-  ]
-
-
 @pytest.mark.asyncio
 async def test_after_model_callback():
   responses = ['model_response']
@@ -136,3 +150,61 @@ async def test_after_model_callback():
   ) == [
       ('root_agent', 'after_model_callback'),
   ]
+
+
+@pytest.mark.asyncio
+async def test_async_before_model_callback():
+  responses = ['model_response']
+  mock_model = utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=MockAsyncBeforeModelCallback(
+          mock_response='async_before_model_callback'
+      ),
+  )
+
+  runner = utils.TestInMemoryRunner(agent)
+  assert utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'async_before_model_callback'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_async_before_model_callback_noop():
+  responses = ['model_response']
+  mock_model = utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      before_model_callback=async_noop_callback,
+  )
+
+  runner = utils.TestInMemoryRunner(agent)
+  assert utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'model_response'),
+  ]
+
+
+@pytest.mark.asyncio
+async def test_async_after_model_callback():
+  responses = ['model_response']
+  mock_model = utils.MockModel.create(responses=responses)
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      after_model_callback=MockAsyncAfterModelCallback(
+          mock_response='async_after_model_callback'
+      ),
+  )
+
+  runner = utils.TestInMemoryRunner(agent)
+  assert utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [
+      ('root_agent', 'async_after_model_callback'),
+  ]