software/src/utils/retry.py

1"""Retry utilities with exponential backoff."""23from __future__ import annotations45import asyncio6import functools7import logging8import random9import time10from dataclasses import dataclass, field11from collections.abc import Awaitable, Coroutine12from typing import Any, Callable, TypeVar1314from src.exceptions import RateLimitError, TranslationError1516logger = logging.getLogger(__name__)1718T = TypeVar("T")192021@dataclass22class RetryConfig:23    """Configuration for retry behavior."""2425    max_attempts: int = 326    base_delay: float = 1.027    max_delay: float = 60.028    exponential_base: float = 2.029    jitter: bool = True30    retryable_exceptions: tuple[type[Exception], ...] = field(31        default_factory=lambda: (RateLimitError, ConnectionError, TimeoutError)32    )3334    def get_delay(self, attempt: int) -> float:35        """Calculate delay for the given attempt number.3637        Uses exponential backoff with optional jitter.3839        Args:40            attempt: The attempt number (0-indexed).4142        Returns:43            Delay in seconds before the next retry.44        """45        delay = self.base_delay * (self.exponential_base**attempt)46        delay = min(delay, self.max_delay)4748        if self.jitter:49            # Add random jitter of +/- 25%50            jitter_range = delay * 0.2551            delay += random.uniform(-jitter_range, jitter_range)5253        return max(0, delay)545556def retry_with_backoff(57    config: RetryConfig | None = None,58    on_retry: Callable[[Exception, int], None] | None = None,59) -> Callable[[Callable[..., T]], Callable[..., T]]:60    """Decorator for retrying functions with exponential backoff.6162    Args:63        config: Retry configuration. Uses defaults if not provided.64        on_retry: Optional callback called before each retry with (exception, attempt).6566    Returns:67        Decorator function.68    """69    if config is None:70        config = RetryConfig()7172    def decorator(func: Callable[..., T]) -> Callable[..., T]:73        @functools.wraps(func)74        def wrapper(*args: Any, **kwargs: Any) -> T:75            last_exception: Exception | None = None7677            for attempt in range(config.max_attempts):78                try:79                    return func(*args, **kwargs)80                except config.retryable_exceptions as e:81                    last_exception = e8283                    # Don't sleep after the last attempt84                    if attempt < config.max_attempts - 1:85                        # Check for rate limit specific retry-after86                        if isinstance(e, RateLimitError) and e.retry_after:87                            delay = e.retry_after88                        else:89                            delay = config.get_delay(attempt)9091                        logger.warning(92                            "Attempt %d/%d failed: %s. Retrying in %.1fs...",93                            attempt + 1,94                            config.max_attempts,95                            str(e),96                            delay,97                        )9899                        if on_retry:100                            on_retry(e, attempt)101102                        time.sleep(delay)103                except Exception:104                    # Non-retryable exception, re-raise immediately105                    raise106107            # All attempts exhausted108            if last_exception:109                raise last_exception110            raise TranslationError("All retry attempts failed")111112        return wrapper113114    return decorator115116117def async_retry_with_backoff(118    config: RetryConfig | None = None,119    on_retry: Callable[[Exception, int], None] | None = None,120) -> Callable[[Callable[..., T]], Callable[..., T]]:121    """Async decorator for retrying functions with exponential backoff.122123    Args:124        config: Retry configuration. Uses defaults if not provided.125        on_retry: Optional callback called before each retry with (exception, attempt).126127    Returns:128        Decorator function for async functions.129    """130    if config is None:131        config = RetryConfig()132133    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Coroutine[Any, Any, T]]:134        @functools.wraps(func)135        async def wrapper(*args: Any, **kwargs: Any) -> T:136            last_exception: Exception | None = None137138            for attempt in range(config.max_attempts):139                try:140                    return await func(*args, **kwargs)141                except config.retryable_exceptions as e:142                    last_exception = e143144                    # Don't sleep after the last attempt145                    if attempt < config.max_attempts - 1:146                        # Check for rate limit specific retry-after147                        if isinstance(e, RateLimitError) and e.retry_after:148                            delay = e.retry_after149                        else:150                            delay = config.get_delay(attempt)151152                        logger.warning(153                            "Attempt %d/%d failed: %s. Retrying in %.1fs...",154                            attempt + 1,155                            config.max_attempts,156                            str(e),157                            delay,158                        )159160                        if on_retry:161                            on_retry(e, attempt)162163                        await asyncio.sleep(delay)164                except Exception:165                    # Non-retryable exception, re-raise immediately166                    raise167168            # All attempts exhausted169            if last_exception:170                raise last_exception171            raise TranslationError("All retry attempts failed")172173        return wrapper174175    return decorator176177178class RetryState:179    """Tracks retry state for manual retry loops."""180181    def __init__(self, config: RetryConfig | None = None) -> None:182        """Initialize retry state.183184        Args:185            config: Retry configuration.186        """187        self.config = config or RetryConfig()188        self.attempt = 0189        self.last_exception: Exception | None = None190191    def should_retry(self, exception: Exception) -> bool:192        """Check if we should retry after an exception.193194        Args:195            exception: The exception that was raised.196197        Returns:198            True if we should retry, False otherwise.199        """200        self.last_exception = exception201202        if not isinstance(exception, self.config.retryable_exceptions):203            return False204205        if self.attempt >= self.config.max_attempts - 1:206            return False207208        return True209210    def wait(self) -> None:211        """Wait before the next retry attempt."""212        if isinstance(self.last_exception, RateLimitError):213            delay = self.last_exception.retry_after or self.config.get_delay(self.attempt)214        else:215            delay = self.config.get_delay(self.attempt)216217        time.sleep(delay)218        self.attempt += 1219220    async def async_wait(self) -> None:221        """Async wait before the next retry attempt."""222        if isinstance(self.last_exception, RateLimitError):223            delay = self.last_exception.retry_after or self.config.get_delay(self.attempt)224        else:225            delay = self.config.get_delay(self.attempt)226227        await asyncio.sleep(delay)228        self.attempt += 1229230    def reset(self) -> None:231        """Reset retry state for a new operation."""232        self.attempt = 0233        self.last_exception = None234