codekingpro/portable-devtools

1"""Model retry middleware for agents."""2 3from __future__ import annotations4 5import asyncio6import time7from typing import TYPE_CHECKING8 9from langchain_core.messages import AIMessage10 11from langchain.agents.middleware._retry import (12    OnFailure,13    RetryOn,14    calculate_delay,15    should_retry_exception,16    validate_retry_params,17)18from langchain.agents.middleware.types import (19    AgentMiddleware,20    AgentState,21    ContextT,22    ModelRequest,23    ModelResponse,24    ResponseT,25)26 27if TYPE_CHECKING:28    from collections.abc import Awaitable, Callable29 30 31class ModelRetryMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]):32    """Middleware that automatically retries failed model calls with configurable backoff.33 34    Supports retrying on specific exceptions and exponential backoff.35 36    Examples:37        !!! example "Basic usage with default settings (2 retries, exponential backoff)"38 39            ```python40            from langchain.agents import create_agent41            from langchain.agents.middleware import ModelRetryMiddleware42 43            agent = create_agent(model, tools=[search_tool], middleware=[ModelRetryMiddleware()])44            ```45 46        !!! example "Retry specific exceptions only"47 48            ```python49            from anthropic import RateLimitError50            from openai import APITimeoutError51 52            retry = ModelRetryMiddleware(53                max_retries=4,54                retry_on=(APITimeoutError, RateLimitError),55                backoff_factor=1.5,56            )57            ```58 59        !!! example "Custom exception filtering"60 61            ```python62            from anthropic import APIStatusError63 64 65            def should_retry(exc: Exception) -> bool:66                # Only retry on 5xx errors67                if isinstance(exc, APIStatusError):68                    return 500 <= exc.status_code < 60069                return False70 71 72            retry = ModelRetryMiddleware(73                max_retries=3,74                retry_on=should_retry,75            )76            ```77 78        !!! example "Custom error handling"79 80            ```python81            def format_error(exc: Exception) -> str:82                return "Model temporarily unavailable. Please try again later."83 84 85            retry = ModelRetryMiddleware(86                max_retries=4,87                on_failure=format_error,88            )89            ```90 91        !!! example "Constant backoff (no exponential growth)"92 93            ```python94            retry = ModelRetryMiddleware(95                max_retries=5,96                backoff_factor=0.0,  # No exponential growth97                initial_delay=2.0,  # Always wait 2 seconds98            )99            ```100 101        !!! example "Raise exception on failure"102 103            ```python104            retry = ModelRetryMiddleware(105                max_retries=2,106                on_failure="error",  # Re-raise exception instead of returning message107            )108            ```109    """110 111    def __init__(112        self,113        *,114        max_retries: int = 2,115        retry_on: RetryOn = (Exception,),116        on_failure: OnFailure = "continue",117        backoff_factor: float = 2.0,118        initial_delay: float = 1.0,119        max_delay: float = 60.0,120        jitter: bool = True,121    ) -> None:122        """Initialize `ModelRetryMiddleware`.123 124        Args:125            max_retries: Maximum number of retry attempts after the initial call.126 127                Must be `>= 0`.128            retry_on: Either a tuple of exception types to retry on, or a callable129                that takes an exception and returns `True` if it should be retried.130 131                Default is to retry on all exceptions.132            on_failure: Behavior when all retries are exhausted.133 134                Options:135 136                - `'continue'`: Return an `AIMessage` with error details,137                    allowing the agent to continue with an error response.138                - `'error'`: Re-raise the exception, stopping agent execution.139                - **Custom callable:** Function that takes the exception and returns a140                    string for the `AIMessage` content, allowing custom error141                    formatting.142            backoff_factor: Multiplier for exponential backoff.143 144                Each retry waits `initial_delay * (backoff_factor ** retry_number)`145                seconds.146 147                Set to `0.0` for constant delay.148            initial_delay: Initial delay in seconds before first retry.149            max_delay: Maximum delay in seconds between retries.150 151                Caps exponential backoff growth.152            jitter: Whether to add random jitter (`±25%`) to delay to avoid thundering herd.153 154        Raises:155            ValueError: If `max_retries < 0` or delays are negative.156        """157        super().__init__()158 159        # Validate parameters160        validate_retry_params(max_retries, initial_delay, max_delay, backoff_factor)161 162        self.max_retries = max_retries163        self.tools = []  # No additional tools registered by this middleware164        self.retry_on = retry_on165        self.on_failure = on_failure166        self.backoff_factor = backoff_factor167        self.initial_delay = initial_delay168        self.max_delay = max_delay169        self.jitter = jitter170 171    @staticmethod172    def _format_failure_message(exc: Exception, attempts_made: int) -> AIMessage:173        """Format the failure message when retries are exhausted.174 175        Args:176            exc: The exception that caused the failure.177            attempts_made: Number of attempts actually made.178 179        Returns:180            `AIMessage` with formatted error message.181        """182        exc_type = type(exc).__name__183        exc_msg = str(exc)184        attempt_word = "attempt" if attempts_made == 1 else "attempts"185        content = (186            f"Model call failed after {attempts_made} {attempt_word} with {exc_type}: {exc_msg}"187        )188        return AIMessage(content=content)189 190    def _handle_failure(self, exc: Exception, attempts_made: int) -> ModelResponse[ResponseT]:191        """Handle failure when all retries are exhausted.192 193        Args:194            exc: The exception that caused the failure.195            attempts_made: Number of attempts actually made.196 197        Returns:198            `ModelResponse` with error details.199 200        Raises:201            Exception: If `on_failure` is `'error'`, re-raises the exception.202        """203        if self.on_failure == "error":204            raise exc205 206        if callable(self.on_failure):207            content = self.on_failure(exc)208            ai_msg = AIMessage(content=content)209        else:210            ai_msg = self._format_failure_message(exc, attempts_made)211 212        return ModelResponse(result=[ai_msg])213 214    def wrap_model_call(215        self,216        request: ModelRequest[ContextT],217        handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],218    ) -> ModelResponse[ResponseT] | AIMessage:219        """Intercept model execution and retry on failure.220 221        Args:222            request: Model request with model, messages, state, and runtime.223            handler: Callable to execute the model (can be called multiple times).224 225        Returns:226            `ModelResponse` or `AIMessage` (the final result).227 228        Raises:229            RuntimeError: If the retry loop completes without returning. (This should not happen.)230        """231        # Initial attempt + retries232        for attempt in range(self.max_retries + 1):233            try:234                return handler(request)235            except Exception as exc:236                attempts_made = attempt + 1  # attempt is 0-indexed237 238                # Check if we should retry this exception239                if not should_retry_exception(exc, self.retry_on):240                    # Exception is not retryable, handle failure immediately241                    return self._handle_failure(exc, attempts_made)242 243                # Check if we have more retries left244                if attempt < self.max_retries:245                    # Calculate and apply backoff delay246                    delay = calculate_delay(247                        attempt,248                        backoff_factor=self.backoff_factor,249                        initial_delay=self.initial_delay,250                        max_delay=self.max_delay,251                        jitter=self.jitter,252                    )253                    if delay > 0:254                        time.sleep(delay)255                    # Continue to next retry256                else:257                    # No more retries, handle failure258                    return self._handle_failure(exc, attempts_made)259 260        # Unreachable: loop always returns via handler success or _handle_failure261        msg = "Unexpected: retry loop completed without returning"262        raise RuntimeError(msg)263 264    async def awrap_model_call(265        self,266        request: ModelRequest[ContextT],267        handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],268    ) -> ModelResponse[ResponseT] | AIMessage:269        """Intercept and control async model execution with retry logic.270 271        Args:272            request: Model request with model, messages, state, and runtime.273            handler: Async callable to execute the model and returns `ModelResponse`.274 275        Returns:276            `ModelResponse` or `AIMessage` (the final result).277 278        Raises:279            RuntimeError: If the retry loop completes without returning. (This should not happen.)280        """281        # Initial attempt + retries282        for attempt in range(self.max_retries + 1):283            try:284                return await handler(request)285            except Exception as exc:286                attempts_made = attempt + 1  # attempt is 0-indexed287 288                # Check if we should retry this exception289                if not should_retry_exception(exc, self.retry_on):290                    # Exception is not retryable, handle failure immediately291                    return self._handle_failure(exc, attempts_made)292 293                # Check if we have more retries left294                if attempt < self.max_retries:295                    # Calculate and apply backoff delay296                    delay = calculate_delay(297                        attempt,298                        backoff_factor=self.backoff_factor,299                        initial_delay=self.initial_delay,300                        max_delay=self.max_delay,301                        jitter=self.jitter,302                    )303                    if delay > 0:304                        await asyncio.sleep(delay)305                    # Continue to next retry306                else:307                    # No more retries, handle failure308                    return self._handle_failure(exc, attempts_made)309 310        # Unreachable: loop always returns via handler success or _handle_failure311        msg = "Unexpected: retry loop completed without returning"312        raise RuntimeError(msg)313