71 lines
1.7 KiB
Python
71 lines
1.7 KiB
Python
"""Common resilience utilities for tool executors."""
|
|
|
|
import asyncio
|
|
import functools
|
|
import logging
|
|
from typing import Optional, Callable, Any, TypeVar
|
|
from tenacity import (
|
|
retry,
|
|
stop_after_attempt,
|
|
wait_exponential,
|
|
retry_if_exception_type,
|
|
before_sleep_log
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Type variable for generic decorators
|
|
T = TypeVar('T')
|
|
|
|
|
|
def async_retry(
|
|
max_attempts: int = 3,
|
|
exceptions: tuple = (Exception,),
|
|
**kwargs
|
|
):
|
|
"""Async retry decorator with exponential backoff.
|
|
|
|
Args:
|
|
max_attempts: Maximum retry attempts
|
|
exceptions: Exception types to retry on
|
|
**kwargs: Additional tenacity configuration
|
|
|
|
Example:
|
|
@async_retry(max_attempts=3, exceptions=(aiohttp.ClientError,))
|
|
async def fetch_data():
|
|
...
|
|
"""
|
|
return retry(
|
|
stop=stop_after_attempt(max_attempts),
|
|
wait=wait_exponential(multiplier=1, min=1, max=10),
|
|
retry=retry_if_exception_type(exceptions),
|
|
reraise=True,
|
|
before_sleep=before_sleep_log(logger, logging.WARNING),
|
|
**kwargs
|
|
)
|
|
|
|
|
|
async def async_timeout_wrapper(
|
|
coro: Callable[..., T],
|
|
timeout: float,
|
|
*args,
|
|
**kwargs
|
|
) -> T:
|
|
"""Wrap async function with timeout.
|
|
|
|
Args:
|
|
coro: Async function to wrap
|
|
timeout: Timeout in seconds
|
|
*args, **kwargs: Arguments for the function
|
|
|
|
Returns:
|
|
Result from the function
|
|
|
|
Raises:
|
|
asyncio.TimeoutError: If timeout exceeded
|
|
|
|
Example:
|
|
result = await async_timeout_wrapper(some_async_func, 5.0, arg1, arg2)
|
|
"""
|
|
return await asyncio.wait_for(coro(*args, **kwargs), timeout=timeout)
|