feat: Add on_model_error_callback in LlmAgent

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

Author: xuanyang15

src/google/adk/agents/llm_agent.py

@@ -85,6 +85,16 @@
     list[_SingleAfterModelCallback],
 ]
 
+_SingleOnModelErrorCallback: TypeAlias = Callable[
+    [CallbackContext, LlmRequest, Exception],
+    Union[Awaitable[Optional[LlmResponse]], Optional[LlmResponse]],
+]
+
+OnModelErrorCallback: TypeAlias = Union[
+    _SingleOnModelErrorCallback,
+    list[_SingleOnModelErrorCallback],
+]
+
 _SingleBeforeToolCallback: TypeAlias = Callable[
     [BaseTool, dict[str, Any], ToolContext],
     Union[Awaitable[Optional[dict]], Optional[dict]],
@@ -364,6 +374,21 @@ class LlmAgent(BaseAgent):
     The content to return to the user. When present, the actual model response
     will be ignored and the provided content will be returned to user.
   """
+  on_model_error_callback: Optional[OnModelErrorCallback] = None
+  """Callback or list of callbacks to be called when a model 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:
+    callback_context: CallbackContext,
+    llm_request: LlmRequest, The raw model request.
+    error: The error from the model call.
+
+  Returns:
+    The content to return to the user. When present, the error will be
+    ignored and the provided content will be returned to user.
+  """
   before_tool_callback: Optional[BeforeToolCallback] = None
   """Callback or list of callbacks to be called before calling the tool.
 
@@ -587,6 +612,20 @@ def canonical_after_model_callbacks(self) -> list[_SingleAfterModelCallback]:
       return self.after_model_callback
     return [self.after_model_callback]
 
+  @property
+  def canonical_on_model_error_callbacks(
+      self,
+  ) -> list[_SingleOnModelErrorCallback]:
+    """The resolved self.on_model_error_callback field as a list of _SingleOnModelErrorCallback.
+
+    This method is only for use by Agent Development Kit.
+    """
+    if not self.on_model_error_callback:
+      return []
+    if isinstance(self.on_model_error_callback, list):
+      return self.on_model_error_callback
+    return [self.on_model_error_callback]
+
   @property
   def canonical_before_tool_callbacks(
       self,

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -977,6 +977,44 @@ async def _run_and_handle_error(
     Yields:
       A generator of LlmResponse.
     """
+
+    from ...agents.llm_agent import LlmAgent
+
+    agent = invocation_context.agent
+    if not isinstance(agent, LlmAgent):
+      raise TypeError(
+          f'Expected agent to be an LlmAgent, but got {type(agent)}'
+      )
+
+    async def _run_on_model_error_callbacks(
+        *,
+        callback_context: CallbackContext,
+        llm_request: LlmRequest,
+        error: Exception,
+    ) -> Optional[LlmResponse]:
+      error_response = (
+          await invocation_context.plugin_manager.run_on_model_error_callback(
+              callback_context=callback_context,
+              llm_request=llm_request,
+              error=error,
+          )
+      )
+      if error_response is not None:
+        return error_response
+
+      for callback in agent.canonical_on_model_error_callbacks:
+        error_response = callback(
+            callback_context=callback_context,
+            llm_request=llm_request,
+            error=error,
+        )
+        if inspect.isawaitable(error_response):
+          error_response = await error_response
+        if error_response is not None:
+          return error_response
+
+      return None
+
     try:
       async with Aclosing(response_generator) as agen:
         async for response in agen:
@@ -985,12 +1023,10 @@ async def _run_and_handle_error(
       callback_context = CallbackContext(
           invocation_context, event_actions=model_response_event.actions
       )
-      error_response = (
-          await invocation_context.plugin_manager.run_on_model_error_callback(
-              callback_context=callback_context,
-              llm_request=llm_request,
-              error=model_error,
-          )
+      error_response = await _run_on_model_error_callbacks(
+          callback_context=callback_context,
+          llm_request=llm_request,
+          error=model_error,
       )
       if error_response is not None:
         yield error_response

tests/unittests/flows/llm_flows/test_model_callbacks.py

@@ -56,6 +56,22 @@ def __call__(
     )
 
 
+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 @@ def supported_models(cls) -> list[str]:
   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 @@ def generate_content(
   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)