Add retry decorator with exponential backoff

examples/retry_decorator_usage.py

@@ -0,0 +1,172 @@
+"""
+Example usage of the retry decorator with exponential backoff.
+
+This example demonstrates how to use the @retry decorator to handle
+flaky API calls and network operations gracefully.
+"""
+
+import random
+from agentunit.core import retry
+
+
+# Example 1: Basic retry with default settings
+@retry()
+def fetch_user_data(user_id: int) -> dict:
+    """
+    Fetch user data from an API with automatic retries.
+
+    Uses default settings:
+    - max_retries=3
+    - base_delay=1.0s
+    - exponential_base=2.0
+    """
+    # Simulated API call that might fail
+    if random.random() < 0.3:  # 30% chance of failure
+        raise ConnectionError("Network timeout")
+
+    return {"id": user_id, "name": "John Doe", "email": "john@example.com"}
+
+
+# Example 2: Custom retry configuration for critical operations
+@retry(
+    max_retries=5,
+    base_delay=2.0,
+    max_delay=30.0,
+    exceptions=(ConnectionError, TimeoutError)
+)
+def critical_database_operation(query: str) -> list:
+    """
+    Execute a critical database query with aggressive retry policy.
+
+    Configuration:
+    - Retries up to 5 times
+    - Starts with 2s delay
+    - Caps delay at 30s
+    - Only retries on connection/timeout errors
+    """
+    # Simulated database operation
+    if random.random() < 0.4:  # 40% chance of failure
+        raise ConnectionError("Database connection lost")
+
+    return [{"result": "data"}]
+
+
+# Example 3: Quick retries for fast operations
+@retry(
+    max_retries=10,
+    base_delay=0.1,
+    max_delay=5.0,
+    exponential_base=1.5
+)
+def quick_cache_lookup(key: str) -> str:
+    """
+    Look up value in cache with quick retries.
+
+    Configuration:
+    - Many retries (10) with short delays
+    - Starts with 0.1s delay
+    - Slower exponential growth (1.5x)
+    - Caps at 5s
+    """
+    # Simulated cache lookup
+    if random.random() < 0.2:  # 20% chance of failure
+        raise ConnectionError("Cache server unavailable")
+
+    return f"value_for_{key}"
+
+
+# Example 4: Specific exception handling
+@retry(
+    max_retries=3,
+    base_delay=1.0,
+    exceptions=(ValueError, KeyError)
+)
+def parse_api_response(response: dict) -> dict:
+    """
+    Parse API response with retries only for specific errors.
+
+    Only retries on ValueError and KeyError.
+    Other exceptions (like TypeError) will raise immediately.
+    """
+    # Simulated parsing that might fail
+    if random.random() < 0.3:
+        raise ValueError("Invalid response format")
+
+    return {"parsed": True, "data": response}
+
+
+# Example 5: Using retry in adapter classes
+class APIAdapter:
+    """Example adapter class using retry decorator."""
+
+    @retry(max_retries=3, base_delay=1.0, exceptions=(ConnectionError,))
+    def call_llm(self, prompt: str) -> str:
+        """
+        Call LLM API with automatic retries on connection errors.
+        """
+        # Simulated LLM API call
+        if random.random() < 0.25:
+            raise ConnectionError("LLM API unavailable")
+
+        return f"Response to: {prompt}"
+
+    @retry(max_retries=5, base_delay=2.0, max_delay=60.0)
+    def fetch_embeddings(self, texts: list[str]) -> list[list[float]]:
+        """
+        Fetch embeddings with more aggressive retry policy.
+        """
+        # Simulated embedding API call
+        if random.random() < 0.3:
+            raise ConnectionError("Embedding service timeout")
+
+        return [[0.1, 0.2, 0.3] for _ in texts]
+
+
+def main():
+    """Demonstrate retry decorator usage."""
+
+    print("Example 1: Basic retry")
+    try:
+        user = fetch_user_data(123)
+        print(f"✓ Fetched user: {user}")
+    except ConnectionError as e:
+        print(f"✗ Failed after retries: {e}")
+
+    print("\nExample 2: Critical operation with custom config")
+    try:
+        results = critical_database_operation("SELECT * FROM users")
+        print(f"✓ Query succeeded: {results}")
+    except ConnectionError as e:
+        print(f"✗ Failed after retries: {e}")
+
+    print("\nExample 3: Quick cache lookup")
+    try:
+        value = quick_cache_lookup("user:123")
+        print(f"✓ Cache hit: {value}")
+    except ConnectionError as e:
+        print(f"✗ Cache unavailable: {e}")
+
+    print("\nExample 4: Specific exception handling")
+    try:
+        parsed = parse_api_response({"status": "ok"})
+        print(f"✓ Parsed response: {parsed}")
+    except ValueError as e:
+        print(f"✗ Parse failed: {e}")
+
+    print("\nExample 5: Adapter class usage")
+    adapter = APIAdapter()
+    try:
+        response = adapter.call_llm("What is AI?")
+        print(f"✓ LLM response: {response}")
+    except ConnectionError as e:
+        print(f"✗ LLM call failed: {e}")
+
+    try:
+        embeddings = adapter.fetch_embeddings(["hello", "world"])
+        print(f"✓ Got embeddings: {len(embeddings)} vectors")
+    except ConnectionError as e:
+        print(f"✗ Embedding fetch failed: {e}")
+
+
+if __name__ == "__main__":
+    main()

src/agentunit/core/__init__.py

@@ -7,6 +7,15 @@
 
 from .runner import Runner, run_suite
 from .scenario import Scenario
+from .utils import retry
 
 
-__all__ = ["DatasetCase", "DatasetSource", "Runner", "Scenario", "ScenarioResult", "run_suite"]
+__all__ = [
+    "DatasetCase",
+    "DatasetSource",
+    "Runner",
+    "Scenario",
+    "ScenarioResult",
+    "retry",
+    "run_suite",
+]

src/agentunit/core/utils.py

@@ -0,0 +1,85 @@
+"""
+Utility functions for AgentUnit core functionality.
+"""
+
+import functools
+import logging
+import time
+from typing import Any, Callable, TypeVar, cast
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def retry(
+    max_retries: int = 3,
+    base_delay: float = 1.0,
+    max_delay: float = 60.0,
+    exponential_base: float = 2.0,
+    exceptions: tuple[type[Exception], ...] = (Exception,),
+) -> Callable[[F], F]:
+    """
+    Decorator that retries a function with exponential backoff on failure.
+
+    This decorator is useful for handling flaky API calls or network operations
+    that may fail temporarily. It implements exponential backoff to avoid
+    overwhelming the service with rapid retry attempts.
+
+    Args:
+        max_retries: Maximum number of retry attempts (default: 3).
+        base_delay: Initial delay between retries in seconds (default: 1.0).
+        max_delay: Maximum delay between retries in seconds (default: 60.0).
+        exponential_base: Base for exponential backoff calculation (default: 2.0).
+        exceptions: Tuple of exception types to catch and retry (default: (Exception,)).
+
+    Returns:
+        Decorated function that implements retry logic with exponential backoff.
+
+    Example:
+        >>> @retry(max_retries=3, base_delay=1.0, exceptions=(ConnectionError,))
+        ... def fetch_data():
+        ...     # API call that might fail
+        ...     return api.get_data()
+
+        >>> @retry(max_retries=5, base_delay=2.0, max_delay=30.0)
+        ... def process_request():
+        ...     # Processing that might need retries
+        ...     return process()
+    """
+
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            last_exception = None
+
+            for attempt in range(max_retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except exceptions as e:
+                    last_exception = e
+
+                    if attempt == max_retries:
+                        logger.error(
+                            f"Function {func.__name__} failed after {max_retries} retries. "
+                            f"Last error: {e}"
+                        )
+                        raise
+
+                    # Calculate delay with exponential backoff
+                    delay = min(base_delay * (exponential_base**attempt), max_delay)
+
+                    logger.warning(
+                        f"Function {func.__name__} failed on attempt {attempt + 1}/{max_retries + 1}. "
+                        f"Retrying in {delay:.2f}s. Error: {e}"
+                    )
+
+                    time.sleep(delay)
+
+            # This should never be reached, but satisfies type checker
+            if last_exception:
+                raise last_exception
+
+        return cast(F, wrapper)
+
+    return decorator