feat: Extend ReflectAndRetryToolPlugin to support hallucinating function calls

PiperOrigin-RevId: 820051762

Author: xuanyang15

contributing/samples/plugin_reflect_tool_retry/README.md

@@ -46,8 +46,30 @@ You can run the agent with:
 $ adk web contributing/samples/plugin_reflect_tool_retry
 ```
 
-You can provide the following prompt to see the agent retrying tool calls:
+Select "basic" and provide the following prompt to see the agent retrying tool
+calls:
 
 ```
 Please guess a number! Tell me what number you guess and how is it.
 ```
+
+### Hallucinating tool calls
+
+The "hallucinating_func_name" agent is an example to show the plugin can retry
+hallucinating tool calls.
+
+For example, we used the `after_model_callback` to hack a tool call with the
+wrong name then the agent can retry calling with the right tool name.
+
+You can run the agent with:
+
+```bash
+$ adk web contributing/samples/plugin_reflect_tool_retry
+```
+
+Select "hallucinating_func_name" and provide the following prompt to see the
+agent retrying tool calls:
+
+```
+Roll a 6 sided die
+```

contributing/samples/plugin_reflect_tool_retry/hallucinating_func_name/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 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 . import agent

contributing/samples/plugin_reflect_tool_retry/hallucinating_func_name/agent.py

@@ -0,0 +1,81 @@
+# Copyright 2025 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.
+
+import random
+
+from google.adk.agents import LlmAgent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.apps.app import App
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins import ReflectAndRetryToolPlugin
+from google.adk.tools.tool_context import ToolContext
+
+APP_NAME = "hallucinating_func_name"
+USER_ID = "test_user"
+
+hallucinated = False  # Whether the tool name is hallucinated
+
+
+def roll_die(sides: int, tool_context: ToolContext) -> int:
+  """Roll a die and return the rolled result.
+
+  Args:
+    sides: The integer number of sides the die has.
+
+  Returns:
+    An integer of the result of rolling the die.
+  """
+  result = random.randint(1, sides)
+  if not "rolls" in tool_context.state:
+    tool_context.state["rolls"] = []
+
+  tool_context.state["rolls"] = tool_context.state["rolls"] + [result]
+  return result
+
+
+def after_model_callback(
+    callback_context: CallbackContext, llm_response: LlmResponse
+):
+  """After model callback to produce one hallucinating tool call."""
+  global hallucinated
+
+  if hallucinated:
+    return None
+
+  if (
+      llm_response.content
+      and llm_response.content.parts[0].function_call.name == "roll_die"
+  ):
+    llm_response.content.parts[0].function_call.name = "roll_die_wrong_name"
+    hallucinated = True
+  return None
+
+
+root_agent = LlmAgent(
+    name="hello_world",
+    description="Helpful agent",
+    instruction="""Use guess_number_tool to guess a number.""",
+    model="gemini-2.5-flash",
+    tools=[roll_die],
+    after_model_callback=after_model_callback,
+)
+
+
+app = App(
+    name=APP_NAME,
+    root_agent=root_agent,
+    plugins=[
+        ReflectAndRetryToolPlugin(max_retries=3),
+    ],
+)

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

@@ -275,21 +275,37 @@ 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."""
-  tool, tool_context = _get_tool_and_context(
-      invocation_context,
-      function_call,
-      tools_dict,
-      tool_confirmation,
+  # 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.
+  function_args = (
+      copy.deepcopy(function_call.args) if function_call.args else {}
   )
 
-  with tracer.start_as_current_span(f'execute_tool {tool.name}'):
-    # 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.
-    function_args = (
-        copy.deepcopy(function_call.args) if function_call.args else {}
+  tool_context = _create_tool_context(
+      invocation_context, function_call, tool_confirmation
+  )
+
+  try:
+    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,
+        )
     )
+    if error_response is not None:
+      return __build_response_event(
+          tool, error_response, tool_context, invocation_context
+      )
+    else:
+      raise tool_error
 
+  with tracer.start_as_current_span(f'execute_tool {tool.name}'):
     # Step 1: Check if plugin before_tool_callback overrides the function
     # response.
     function_response = (
@@ -639,24 +655,45 @@ async def run_tool_and_update_queue(tool, function_args, tool_context):
   return function_response
 
 
-def _get_tool_and_context(
-    invocation_context: InvocationContext,
-    function_call: types.FunctionCall,
-    tools_dict: dict[str, BaseTool],
-    tool_confirmation: Optional[ToolConfirmation] = None,
+def _get_tool(
+    function_call: types.FunctionCall, tools_dict: dict[str, BaseTool]
 ):
+  """Returns the tool corresponding to the function call."""
   if function_call.name not in tools_dict:
     raise ValueError(
-        f'Function {function_call.name} is not found in the tools_dict.'
+        f'Function {function_call.name} is not found in the tools_dict:'
+        f' {tools_dict.keys()}.'
     )
 
-  tool_context = ToolContext(
+  return tools_dict[function_call.name]
+
+
+def _create_tool_context(
+    invocation_context: InvocationContext,
+    function_call: types.FunctionCall,
+    tool_confirmation: Optional[ToolConfirmation] = None,
+):
+  """Creates a ToolContext object."""
+  return ToolContext(
       invocation_context=invocation_context,
       function_call_id=function_call.id,
       tool_confirmation=tool_confirmation,
   )
 
-  tool = tools_dict[function_call.name]
+
+def _get_tool_and_context(
+    invocation_context: InvocationContext,
+    function_call: types.FunctionCall,
+    tools_dict: dict[str, BaseTool],
+    tool_confirmation: Optional[ToolConfirmation] = None,
+):
+  """Returns the tool and tool context corresponding to the function call."""
+  tool = _get_tool(function_call, tools_dict)
+  tool_context = _create_tool_context(
+      invocation_context,
+      function_call,
+      tool_confirmation,
+  )
 
   return (tool, tool_context)
 

src/google/adk/plugins/reflect_retry_tool_plugin.py

@@ -142,8 +142,20 @@ async def after_tool_callback(
       tool_args: dict[str, Any],
       tool_context: ToolContext,
       result: Any,
-  ) -> Optional[dict]:
-    """Handles successful tool calls or extracts and processes errors."""
+  ) -> Optional[dict[str, Any]]:
+    """Handles successful tool calls or extracts and processes errors.
+
+    Args:
+      tool: The tool that was called.
+      tool_args: The arguments passed to the tool.
+      tool_context: The context of the tool call.
+      result: The result of the tool call.
+
+    Returns:
+      An optional dictionary containing reflection guidance if an error is
+      detected, or None if the tool call was successful or the
+      response is already a reflection message.
+    """
     if (
         isinstance(result, dict)
         and result.get("response_type") == REFLECT_AND_RETRY_RESPONSE_TYPE
@@ -157,7 +169,8 @@ async def after_tool_callback(
     if error:
       return await self._handle_tool_error(tool, tool_args, tool_context, error)
 
-    # On success, reset the failure count for this specific tool within its scope.
+    # On success, reset the failure count for this specific tool within
+    # its scope.
     await self._reset_failures_for_tool(tool_context, tool.name)
     return None
 
@@ -168,14 +181,23 @@ async def extract_error_from_result(
       tool_args: dict[str, Any],
       tool_context: ToolContext,
       result: Any,
-  ) -> Optional[Any]:
+  ) -> Optional[dict[str, Any]]:
     """Extracts an error from a successful tool result and triggers retry logic.
 
     This is useful when tool call finishes successfully but the result contains
     an error object like {"error": ...} that should be handled by the plugin.
 
     By overriding this method, you can trigger retry logic on these successful
     results that contain errors.
+
+    Args:
+      tool: The tool that was called.
+      tool_args: The arguments passed to the tool.
+      tool_context: The context of the tool call.
+      result: The result of the tool call.
+
+    Returns:
+      The extracted error if any, or None if no error was detected.
     """
     return None
 
@@ -186,8 +208,18 @@ async def on_tool_error_callback(
       tool_args: dict[str, Any],
       tool_context: ToolContext,
       error: Exception,
-  ) -> Optional[dict]:
-    """Handles tool exceptions by providing reflection guidance."""
+  ) -> Optional[dict[str, Any]]:
+    """Handles tool exceptions by providing reflection guidance.
+
+    Args:
+      tool: The tool that was called.
+      tool_args: The arguments passed to the tool.
+      tool_context: The context of the tool call.
+      error: The exception raised by the tool.
+
+    Returns:
+      An optional dictionary containing reflection guidance for the error.
+    """
     return await self._handle_tool_error(tool, tool_args, tool_context, error)
 
   async def _handle_tool_error(
@@ -196,8 +228,18 @@ async def _handle_tool_error(
       tool_args: dict[str, Any],
       tool_context: ToolContext,
       error: Any,
-  ) -> Optional[dict]:
-    """Central, thread-safe logic for processing tool errors."""
+  ) -> Optional[dict[str, Any]]:
+    """Central, thread-safe logic for processing tool errors.
+
+    Args:
+      tool: The tool that was called.
+      tool_args: The arguments passed to the tool.
+      tool_context: The context of the tool call.
+      error: The error to be handled.
+
+    Returns:
+      An optional dictionary containing reflection guidance for the error.
+    """
     if self.max_retries == 0:
       if self.throw_exception_if_retry_exceeded:
         raise error
@@ -285,6 +327,7 @@ def _create_tool_reflection_response(
 2.  **State or Preconditions**: Did a previous step fail or not produce the necessary state/resource for this tool to succeed?
 3.  **Alternative Approach**: Is this the right tool for the job? Could another tool or a different sequence of steps achieve the goal?
 4.  **Simplify the Task**: Can you break the problem down into smaller, simpler steps?
+5.  **Wrong Function Name**: Does the error indicates the tool is not found? Please check again and only use available tools.
 
 Formulate a new plan based on your analysis and try a corrected or different approach.
 """