Skip to content

Commit e55b894

Browse files
pwwpchecopybara-github
authored andcommitted
feat: Add ReflectRetryToolPlugin to reflect from errors and retry with different arguments when tool errors
This plugin intercepts tool failures, provides structured guidance to the LLM for reflection and correction, and retries the operation up to a configurable limit. **Key Features:** - **Concurrency Safe:** Uses locking to safely handle parallel tool executions - **Configurable Scope:** Tracks failures per-invocation (default) or globally using the `TrackingScope` enum. - **Extensible Scoping:** The `_get_scope_key` method can be overridden to implement custom tracking logic (e.g., per-user or per-session). - **Granular Tracking:** Failure counts are tracked per-tool within the defined scope. A success with one tool resets its counter without affecting others. - **Custom Error Extraction:** Supports detecting errors in normal tool responses that don't throw exceptions, by overriding the `extract_error_from_result` method. **Example:** ```python from my_project.plugins import ReflectAndRetryToolPlugin, TrackingScope # Example 1: (MOST COMMON USAGE): # Track failures only within the current agent invocation (default). error_handling_plugin = ReflectAndRetryToolPlugin(max_retries=3) # Example 2: # Track failures globally across all turns and users. global_error_handling_plugin = ReflectAndRetryToolPlugin(max_retries=5, scope=TrackingScope.GLOBAL) # Example 3: # Retry on failures but do not throw exceptions. error_handling_plugin = ReflectAndRetryToolPlugin(max_retries=3, throw_exception_if_retry_exceeded=False) # Example 4: # Track failures in successful tool responses that contain errors. class CustomRetryPlugin(ReflectAndRetryToolPlugin): async def extract_error_from_result(self, *, tool, tool_args,tool_context, result): # Detect error based on response content if result.get('status') == 'error': return result return None # No error detected error_handling_plugin = CustomRetryPlugin(max_retries=5) ``` PiperOrigin-RevId: 816456549
1 parent 2b5acb9 commit e55b894

4 files changed

Lines changed: 878 additions & 2 deletions

File tree

src/google/adk/plugins/__init__.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,5 +13,13 @@
1313
# limitations under the License.
1414

1515
from .base_plugin import BasePlugin
16+
from .logging_plugin import LoggingPlugin
17+
from .plugin_manager import PluginManager
18+
from .reflect_retry_tool_plugin import ReflectAndRetryToolPlugin
1619

17-
__all__ = ['BasePlugin']
20+
__all__ = [
21+
'BasePlugin',
22+
'LoggingPlugin',
23+
'PluginManager',
24+
'ReflectAndRetryToolPlugin',
25+
]

src/google/adk/plugins/logging_plugin.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,22 @@
1616

1717
from typing import Any
1818
from typing import Optional
19+
from typing import TYPE_CHECKING
1920

2021
from google.genai import types
2122

2223
from ..agents.base_agent import BaseAgent
2324
from ..agents.callback_context import CallbackContext
24-
from ..agents.invocation_context import InvocationContext
2525
from ..events.event import Event
2626
from ..models.llm_request import LlmRequest
2727
from ..models.llm_response import LlmResponse
2828
from ..tools.base_tool import BaseTool
2929
from ..tools.tool_context import ToolContext
3030
from .base_plugin import BasePlugin
3131

32+
if TYPE_CHECKING:
33+
from ..agents.invocation_context import InvocationContext
34+
3235

3336
class LoggingPlugin(BasePlugin):
3437
"""A plugin that logs important information at each callback point.
Lines changed: 339 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,339 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import asyncio
18+
from enum import Enum
19+
import json
20+
from typing import Any
21+
from typing import Optional
22+
23+
from pydantic import BaseModel
24+
25+
from ..tools.base_tool import BaseTool
26+
from ..tools.tool_context import ToolContext
27+
from ..utils.feature_decorator import experimental
28+
from .base_plugin import BasePlugin
29+
30+
REFLECT_AND_RETRY_RESPONSE_TYPE = "ERROR_HANDLED_BY_REFLECT_AND_RETRY_PLUGIN"
31+
GLOBAL_SCOPE_KEY = "__global_reflect_and_retry_scope__"
32+
33+
# A mapping from a tool's name to its consecutive failure count.
34+
PerToolFailuresCounter = dict[str, int]
35+
36+
37+
class TrackingScope(Enum):
38+
"""Defines the lifecycle scope for tracking tool failure counts."""
39+
40+
INVOCATION = "invocation"
41+
GLOBAL = "global"
42+
43+
44+
class ToolFailureResponse(BaseModel):
45+
"""Response containing tool failure details and retry guidance."""
46+
47+
response_type: str = REFLECT_AND_RETRY_RESPONSE_TYPE
48+
error_type: str = ""
49+
error_details: str = ""
50+
retry_count: int = 0
51+
reflection_guidance: str = ""
52+
53+
54+
@experimental
55+
class ReflectAndRetryToolPlugin(BasePlugin):
56+
"""Provides self-healing, concurrent-safe error recovery for tool failures.
57+
58+
This plugin intercepts tool failures, provides structured guidance to the LLM
59+
for reflection and correction, and retries the operation up to a configurable
60+
limit.
61+
62+
**Key Features:**
63+
64+
- **Concurrency Safe:** Uses locking to safely handle parallel tool
65+
executions
66+
- **Configurable Scope:** Tracks failures per-invocation (default) or globally
67+
using the `TrackingScope` enum.
68+
- **Extensible Scoping:** The `_get_scope_key` method can be overridden to
69+
implement custom tracking logic (e.g., per-user or per-session).
70+
- **Granular Tracking:** Failure counts are tracked per-tool within the
71+
defined scope. A success with one tool resets its counter without affecting
72+
others.
73+
- **Custom Error Extraction:** Supports detecting errors in normal tool
74+
responses
75+
that
76+
don't throw exceptions, by overriding the `extract_error_from_result`
77+
method.
78+
79+
**Example:**
80+
```python
81+
from my_project.plugins import ReflectAndRetryToolPlugin, TrackingScope
82+
83+
# Example 1: (MOST COMMON USAGE):
84+
# Track failures only within the current agent invocation (default).
85+
error_handling_plugin = ReflectAndRetryToolPlugin(max_retries=3)
86+
87+
# Example 2:
88+
# Track failures globally across all turns and users.
89+
global_error_handling_plugin = ReflectAndRetryToolPlugin(max_retries=5,
90+
scope=TrackingScope.GLOBAL)
91+
92+
# Example 3:
93+
# Retry on failures but do not throw exceptions.
94+
error_handling_plugin =
95+
ReflectAndRetryToolPlugin(max_retries=3,
96+
throw_exception_if_retry_exceeded=False)
97+
98+
# Example 4:
99+
# Track failures in successful tool responses that contain errors.
100+
class CustomRetryPlugin(ReflectAndRetryToolPlugin):
101+
async def extract_error_from_result(self, *, tool, tool_args,tool_context,
102+
result):
103+
# Detect error based on response content
104+
if result.get('status') == 'error':
105+
return result
106+
return None # No error detected
107+
error_handling_plugin = CustomRetryPlugin(max_retries=5)
108+
```
109+
"""
110+
111+
def __init__(
112+
self,
113+
name: str = "reflect_retry_tool_plugin",
114+
max_retries: int = 3,
115+
throw_exception_if_retry_exceeded: bool = True,
116+
tracking_scope: TrackingScope = TrackingScope.INVOCATION,
117+
):
118+
"""Initializes the ReflectAndRetryToolPlugin.
119+
120+
Args:
121+
name: Plugin instance identifier.
122+
max_retries: Maximum consecutive failures before giving up (0 = no
123+
retries).
124+
throw_exception_if_retry_exceeded: If True, raises the final exception
125+
when the retry limit is reached. If False, returns guidance instead.
126+
tracking_scope: Determines the lifecycle of the error tracking state.
127+
Defaults to `TrackingScope.INVOCATION` tracking per-invocation.
128+
"""
129+
super().__init__(name)
130+
if max_retries < 0:
131+
raise ValueError("max_retries must be a non-negative integer.")
132+
self.max_retries = max_retries
133+
self.throw_exception_if_retry_exceeded = throw_exception_if_retry_exceeded
134+
self.scope = tracking_scope
135+
self._scoped_failure_counters: dict[str, PerToolFailuresCounter] = {}
136+
self._lock = asyncio.Lock()
137+
138+
async def after_tool_callback(
139+
self,
140+
*,
141+
tool: BaseTool,
142+
tool_args: dict[str, Any],
143+
tool_context: ToolContext,
144+
result: Any,
145+
) -> Optional[dict]:
146+
"""Handles successful tool calls or extracts and processes errors."""
147+
if (
148+
isinstance(result, dict)
149+
and result.get("response_type") == REFLECT_AND_RETRY_RESPONSE_TYPE
150+
):
151+
return None
152+
153+
error = await self.extract_error_from_result(
154+
tool=tool, tool_args=tool_args, tool_context=tool_context, result=result
155+
)
156+
157+
if error:
158+
return await self._handle_tool_error(tool, tool_args, tool_context, error)
159+
160+
# On success, reset the failure count for this specific tool within its scope.
161+
await self._reset_failures_for_tool(tool_context, tool.name)
162+
return None
163+
164+
async def extract_error_from_result(
165+
self,
166+
*,
167+
tool: BaseTool,
168+
tool_args: dict[str, Any],
169+
tool_context: ToolContext,
170+
result: Any,
171+
) -> Optional[Any]:
172+
"""Extracts an error from a successful tool result and triggers retry logic.
173+
174+
This is useful when tool call finishes successfully but the result contains
175+
an error object like {"error": ...} that should be handled by the plugin.
176+
177+
By overriding this method, you can trigger retry logic on these successful
178+
results that contain errors.
179+
"""
180+
return None
181+
182+
async def on_tool_error_callback(
183+
self,
184+
*,
185+
tool: BaseTool,
186+
tool_args: dict[str, Any],
187+
tool_context: ToolContext,
188+
error: Exception,
189+
) -> Optional[dict]:
190+
"""Handles tool exceptions by providing reflection guidance."""
191+
return await self._handle_tool_error(tool, tool_args, tool_context, error)
192+
193+
async def _handle_tool_error(
194+
self,
195+
tool: BaseTool,
196+
tool_args: dict[str, Any],
197+
tool_context: ToolContext,
198+
error: Any,
199+
) -> Optional[dict]:
200+
"""Central, thread-safe logic for processing tool errors."""
201+
if self.max_retries == 0:
202+
if self.throw_exception_if_retry_exceeded:
203+
raise error
204+
return self._get_tool_retry_exceed_msg(tool, error, tool_args)
205+
206+
scope_key = self._get_scope_key(tool_context)
207+
async with self._lock:
208+
tool_failure_counter = self._scoped_failure_counters.setdefault(
209+
scope_key, {}
210+
)
211+
current_retries = tool_failure_counter.get(tool.name, 0) + 1
212+
tool_failure_counter[tool.name] = current_retries
213+
214+
if current_retries <= self.max_retries:
215+
return self._create_tool_reflection_response(
216+
tool, tool_args, error, current_retries
217+
)
218+
219+
# Max Retry exceeded
220+
if self.throw_exception_if_retry_exceeded:
221+
raise error
222+
else:
223+
return self._get_tool_retry_exceed_msg(tool, tool_args, error)
224+
225+
def _get_scope_key(self, tool_context: ToolContext) -> str:
226+
"""Returns a unique key for the state dictionary based on the scope.
227+
228+
This method can be overridden in a subclass to implement custom scoping
229+
logic, for example, tracking failures on a per-user or per-session basis.
230+
"""
231+
if self.scope is TrackingScope.INVOCATION:
232+
return tool_context.invocation_id
233+
elif self.scope is TrackingScope.GLOBAL:
234+
return GLOBAL_SCOPE_KEY
235+
raise ValueError(f"Unknown scope: {self.scope}")
236+
237+
async def _reset_failures_for_tool(
238+
self, tool_context: ToolContext, tool_name: str
239+
) -> None:
240+
"""Atomically resets the failure count for a tool and cleans up state."""
241+
scope = self._get_scope_key(tool_context)
242+
async with self._lock:
243+
if scope in self._scoped_failure_counters:
244+
state = self._scoped_failure_counters[scope]
245+
state.pop(tool_name, None)
246+
247+
def _ensure_exception(self, error: Any) -> Exception:
248+
"""Ensures the given error is an Exception instance, wrapping if not."""
249+
return error if isinstance(error, Exception) else Exception(str(error))
250+
251+
def _format_error_details(self, error: Any) -> str:
252+
"""Formats error details for inclusion in the reflection message."""
253+
if isinstance(error, Exception):
254+
return f"{type(error).__name__}: {str(error)}"
255+
return str(error)
256+
257+
def _create_tool_reflection_response(
258+
self,
259+
tool: BaseTool,
260+
tool_args: dict[str, Any],
261+
error: Any,
262+
retry_count: int,
263+
) -> dict[str, Any]:
264+
"""Generates structured reflection guidance for tool failures."""
265+
args_summary = json.dumps(tool_args, indent=2, default=str)
266+
error_details = self._format_error_details(error)
267+
268+
reflection_message = f"""
269+
The call to tool `{tool.name}` failed.
270+
271+
**Error Details:**
272+
```
273+
{error_details}
274+
```
275+
276+
**Tool Arguments Used:**
277+
```json
278+
{args_summary}
279+
```
280+
281+
**Reflection Guidance:**
282+
This is retry attempt **{retry_count} of {self.max_retries}**. Analyze the error and the arguments you provided. Do not repeat the exact same call. Consider the following before your next attempt:
283+
284+
1. **Invalid Parameters**: Does the error suggest that one or more arguments are incorrect, badly formatted, or missing? Review the tool's schema and your arguments.
285+
2. **State or Preconditions**: Did a previous step fail or not produce the necessary state/resource for this tool to succeed?
286+
3. **Alternative Approach**: Is this the right tool for the job? Could another tool or a different sequence of steps achieve the goal?
287+
4. **Simplify the Task**: Can you break the problem down into smaller, simpler steps?
288+
289+
Formulate a new plan based on your analysis and try a corrected or different approach.
290+
"""
291+
292+
return ToolFailureResponse(
293+
error_type=(
294+
type(error).__name__
295+
if isinstance(error, Exception)
296+
else "ToolError"
297+
),
298+
error_details=str(error),
299+
retry_count=retry_count,
300+
reflection_guidance=reflection_message.strip(),
301+
).model_dump(mode="json")
302+
303+
def _get_tool_retry_exceed_msg(
304+
self,
305+
tool: BaseTool,
306+
tool_args: dict[str, Any],
307+
error: Exception,
308+
) -> dict[str, Any]:
309+
"""Generates guidance when the maximum retry limit is exceeded."""
310+
error_details = self._format_error_details(error)
311+
args_summary = json.dumps(tool_args, indent=2, default=str)
312+
313+
reflection_message = f"""
314+
The tool `{tool.name}` has failed consecutively {self.max_retries} times and the retry limit has been exceeded.
315+
316+
**Last Error:**
317+
```
318+
{error_details}
319+
```
320+
321+
**Last Arguments Used:**
322+
```json
323+
{args_summary}
324+
```
325+
326+
**Final Instruction:**
327+
**Do not attempt to use the `{tool.name}` tool again for this task.** You must now try a different approach. Acknowledge the failure and devise a new strategy, potentially using other available tools or informing the user that the task cannot be completed.
328+
"""
329+
330+
return ToolFailureResponse(
331+
error_type=(
332+
type(error).__name__
333+
if isinstance(error, Exception)
334+
else "ToolError"
335+
),
336+
error_details=str(error),
337+
retry_count=self.max_retries,
338+
reflection_guidance=reflection_message.strip(),
339+
).model_dump(mode="json")

0 commit comments

Comments
 (0)