feat: Add on_model_error_callback in LlmAgent

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 828560608

tests/unittests/flows/llm_flows/test_model_callbacks.py

@@ -56,6 +56,22 @@ class MockAfterModelCallback(BaseModel):
     )
 
 
+class MockOnModelCallback(BaseModel):
+  mock_response: str
+
+  def __call__(
+      self,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+      error: Exception,
+  ) -> LlmResponse:
+    return LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.mock_response)]
+        )
+    )
+
+
 def noop_callback(**kwargs) -> Optional[LlmResponse]:
   pass
 
@@ -140,3 +156,40 @@ async def test_after_model_callback_noop():
   assert testing_utils.simplify_events(
       await runner.run_async_with_new_session('test')
   ) == [('root_agent', 'model_response')]
+
+
+@pytest.mark.asyncio
+async def test_on_model_callback_model_error_noop():
+  """Test that the on_model_error_callback is a no-op when the model returns an error."""
+  mock_model = testing_utils.MockModel.create(
+      responses=[], error=SystemError('error')
+  )
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_model_error_callback=noop_callback,
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  with pytest.raises(SystemError):
+    await runner.run_async_with_new_session('test')
+
+
+@pytest.mark.asyncio
+async def test_on_model_callback_model_error_modify_model_response():
+  """Test that the on_model_error_callback can modify the model response."""
+  mock_model = testing_utils.MockModel.create(
+      responses=[], error=SystemError('error')
+  )
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      on_model_error_callback=MockOnModelCallback(
+          mock_response='on_model_error_callback_response'
+      ),
+  )
+
+  runner = testing_utils.TestInMemoryRunner(agent)
+  assert testing_utils.simplify_events(
+      await runner.run_async_with_new_session('test')
+  ) == [('root_agent', 'on_model_error_callback_response')]

tests/unittests/testing_utils.py

@@ -363,7 +363,7 @@ class MockModel(BaseLlm):
   def generate_content(
       self, llm_request: LlmRequest, stream: bool = False
   ) -> Generator[LlmResponse, None, None]:
-    if self.error:
+    if self.error is not None:
       raise self.error
     # Increasement of the index has to happen before the yield.
     self.response_index += 1
@@ -375,6 +375,8 @@ class MockModel(BaseLlm):
   async def generate_content_async(
       self, llm_request: LlmRequest, stream: bool = False
   ) -> AsyncGenerator[LlmResponse, None]:
+    if self.error is not None:
+      raise self.error
     # Increasement of the index has to happen before the yield.
     self.response_index += 1
     self.requests.append(llm_request)