|
| 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