tool improvment

cortex/autonomy/tools/utils/resilience.py

@@ -0,0 +1,70 @@
+"""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)