codekingpro/portable-devtools

1"""Shared retry utilities for agent middleware.2 3This module contains common constants, utilities, and logic used by both4model and tool retry middleware implementations.5"""6 7from __future__ import annotations8 9import random10from collections.abc import Callable11from typing import Literal12 13# Type aliases14RetryOn = tuple[type[Exception], ...] | Callable[[Exception], bool]15"""Type for specifying which exceptions to retry on.16 17Can be either:18- A tuple of exception types to retry on (based on `isinstance` checks)19- A callable that takes an exception and returns `True` if it should be retried20"""21 22OnFailure = Literal["error", "continue"] | Callable[[Exception], str]23"""Type for specifying failure handling behavior.24 25Can be either:26- A literal action string (`'error'` or `'continue'`)27    - `'error'`: Re-raise the exception, stopping agent execution.28    - `'continue'`: Inject a message with the error details, allowing the agent to continue.29       For tool retries, a `ToolMessage` with the error details will be injected.30       For model retries, an `AIMessage` with the error details will be returned.31- A callable that takes an exception and returns a string for error message content32"""33 34 35def validate_retry_params(36    max_retries: int,37    initial_delay: float,38    max_delay: float,39    backoff_factor: float,40) -> None:41    """Validate retry parameters.42 43    Args:44        max_retries: Maximum number of retry attempts.45        initial_delay: Initial delay in seconds before first retry.46        max_delay: Maximum delay in seconds between retries.47        backoff_factor: Multiplier for exponential backoff.48 49    Raises:50        ValueError: If any parameter is invalid (negative values).51    """52    if max_retries < 0:53        msg = "max_retries must be >= 0"54        raise ValueError(msg)55    if initial_delay < 0:56        msg = "initial_delay must be >= 0"57        raise ValueError(msg)58    if max_delay < 0:59        msg = "max_delay must be >= 0"60        raise ValueError(msg)61    if backoff_factor < 0:62        msg = "backoff_factor must be >= 0"63        raise ValueError(msg)64 65 66def should_retry_exception(67    exc: Exception,68    retry_on: RetryOn,69) -> bool:70    """Check if an exception should trigger a retry.71 72    Args:73        exc: The exception that occurred.74        retry_on: Either a tuple of exception types to retry on, or a callable75            that takes an exception and returns `True` if it should be retried.76 77    Returns:78        `True` if the exception should be retried, `False` otherwise.79    """80    if callable(retry_on):81        return retry_on(exc)82    return isinstance(exc, retry_on)83 84 85def calculate_delay(86    retry_number: int,87    *,88    backoff_factor: float,89    initial_delay: float,90    max_delay: float,91    jitter: bool,92) -> float:93    """Calculate delay for a retry attempt with exponential backoff and optional jitter.94 95    Args:96        retry_number: The retry attempt number (0-indexed).97        backoff_factor: Multiplier for exponential backoff.98 99            Set to `0.0` for constant delay.100        initial_delay: Initial delay in seconds before first retry.101        max_delay: Maximum delay in seconds between retries.102 103            Caps exponential backoff growth.104        jitter: Whether to add random jitter to delay to avoid thundering herd.105 106    Returns:107        Delay in seconds before next retry.108    """109    if backoff_factor == 0.0:110        delay = initial_delay111    else:112        delay = initial_delay * (backoff_factor**retry_number)113 114    # Cap at max_delay115    delay = min(delay, max_delay)116 117    if jitter and delay > 0:118        jitter_amount = delay * 0.25  # ±25% jitter119        delay += random.uniform(-jitter_amount, jitter_amount)  # noqa: S311120        # Ensure delay is not negative after jitter121        delay = max(0, delay)122 123    return delay124