agent-tool-resilience 0.1.0__py3-none-any.whl

agent_tool_resilience/rate_limit.py

@@ -0,0 +1,261 @@
+"""
+Rate limit handling for API calls.
+"""
+
+import asyncio
+import threading
+import time
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional
+
+
+class RateLimitExceeded(Exception):
+    """Raised when rate limit is exceeded and blocking is disabled."""
+    
+    def __init__(self, message: str, retry_after: Optional[float] = None):
+        super().__init__(message)
+        self.retry_after = retry_after
+
+
+@dataclass
+class RateLimitHandler:
+    """
+    Handles rate limiting for API calls.
+    
+    Supports:
+    - Request counting with sliding window
+    - Automatic throttling
+    - Respecting Retry-After headers
+    - Multiple rate limit tiers
+    
+    Attributes:
+        requests_per_minute: Maximum requests per minute (0 = unlimited)
+        requests_per_second: Maximum requests per second (0 = unlimited)
+        auto_throttle: Whether to automatically delay requests
+        respect_retry_after: Whether to respect Retry-After from exceptions
+        on_throttle: Callback when throttling occurs
+    """
+    requests_per_minute: int = 0
+    requests_per_second: int = 0
+    auto_throttle: bool = True
+    respect_retry_after: bool = True
+    on_throttle: Optional[Callable[[float], None]] = None
+    
+    _minute_window: deque = field(default_factory=lambda: deque(), init=False)
+    _second_window: deque = field(default_factory=lambda: deque(), init=False)
+    _lock: threading.Lock = field(default_factory=threading.Lock, init=False)
+    _retry_after_until: Optional[float] = field(default=None, init=False)
+    
+    def _clean_windows(self, now: float) -> None:
+        """Remove old entries from sliding windows."""
+        # Clean minute window
+        while self._minute_window and now - self._minute_window[0] > 60:
+            self._minute_window.popleft()
+        
+        # Clean second window
+        while self._second_window and now - self._second_window[0] > 1:
+            self._second_window.popleft()
+    
+    def _get_wait_time(self, now: float) -> float:
+        """Calculate how long to wait before the next request."""
+        wait_time = 0.0
+        
+        # Check retry-after
+        if self._retry_after_until and now < self._retry_after_until:
+            wait_time = max(wait_time, self._retry_after_until - now)
+        
+        # Check per-minute limit
+        if self.requests_per_minute > 0:
+            if len(self._minute_window) >= self.requests_per_minute:
+                oldest = self._minute_window[0]
+                wait_time = max(wait_time, 60 - (now - oldest))
+        
+        # Check per-second limit
+        if self.requests_per_second > 0:
+            if len(self._second_window) >= self.requests_per_second:
+                oldest = self._second_window[0]
+                wait_time = max(wait_time, 1 - (now - oldest))
+        
+        return wait_time
+    
+    def _record_request(self, now: float) -> None:
+        """Record a request in the sliding windows."""
+        if self.requests_per_minute > 0:
+            self._minute_window.append(now)
+        if self.requests_per_second > 0:
+            self._second_window.append(now)
+    
+    def set_retry_after(self, seconds: float) -> None:
+        """Set a retry-after delay from an API response."""
+        with self._lock:
+            self._retry_after_until = time.time() + seconds
+    
+    def check_rate_limit(self) -> float:
+        """
+        Check if we're within rate limits.
+        
+        Returns:
+            Wait time in seconds (0 if no wait needed)
+        """
+        with self._lock:
+            now = time.time()
+            self._clean_windows(now)
+            return self._get_wait_time(now)
+    
+    def acquire(self) -> None:
+        """
+        Acquire permission to make a request.
+        
+        Blocks if auto_throttle is True, otherwise raises RateLimitExceeded.
+        """
+        with self._lock:
+            now = time.time()
+            self._clean_windows(now)
+            wait_time = self._get_wait_time(now)
+            
+            if wait_time > 0:
+                if self.auto_throttle:
+                    if self.on_throttle:
+                        self.on_throttle(wait_time)
+                    time.sleep(wait_time)
+                    now = time.time()
+                else:
+                    raise RateLimitExceeded(
+                        f"Rate limit exceeded, retry after {wait_time:.2f}s",
+                        retry_after=wait_time
+                    )
+            
+            self._record_request(now)
+    
+    async def acquire_async(self) -> None:
+        """
+        Async version of acquire.
+        
+        Blocks if auto_throttle is True, otherwise raises RateLimitExceeded.
+        """
+        with self._lock:
+            now = time.time()
+            self._clean_windows(now)
+            wait_time = self._get_wait_time(now)
+        
+        if wait_time > 0:
+            if self.auto_throttle:
+                if self.on_throttle:
+                    self.on_throttle(wait_time)
+                await asyncio.sleep(wait_time)
+                now = time.time()
+            else:
+                raise RateLimitExceeded(
+                    f"Rate limit exceeded, retry after {wait_time:.2f}s",
+                    retry_after=wait_time
+                )
+        
+        with self._lock:
+            self._record_request(time.time())
+    
+    def execute(
+        self,
+        func: Callable[..., Any],
+        *args: Any,
+        **kwargs: Any
+    ) -> Any:
+        """
+        Execute a function with rate limiting.
+        
+        Args:
+            func: Function to execute
+            *args: Positional arguments
+            **kwargs: Keyword arguments
+            
+        Returns:
+            Function's return value
+        """
+        self.acquire()
+        
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            # Check for rate limit response
+            if self.respect_retry_after:
+                retry_after = self._extract_retry_after(e)
+                if retry_after:
+                    self.set_retry_after(retry_after)
+            raise
+    
+    async def execute_async(
+        self,
+        func: Callable[..., Any],
+        *args: Any,
+        **kwargs: Any
+    ) -> Any:
+        """
+        Execute an async function with rate limiting.
+        
+        Args:
+            func: Async function to execute
+            *args: Positional arguments
+            **kwargs: Keyword arguments
+            
+        Returns:
+            Function's return value
+        """
+        await self.acquire_async()
+        
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            # Check for rate limit response
+            if self.respect_retry_after:
+                retry_after = self._extract_retry_after(e)
+                if retry_after:
+                    self.set_retry_after(retry_after)
+            raise
+    
+    def _extract_retry_after(self, exception: Exception) -> Optional[float]:
+        """Try to extract retry-after from an exception."""
+        # Check for retry_after attribute
+        if hasattr(exception, "retry_after"):
+            return float(exception.retry_after)
+        
+        # Check for response with Retry-After header
+        if hasattr(exception, "response"):
+            response = exception.response
+            if hasattr(response, "headers"):
+                retry_after = response.headers.get("Retry-After")
+                if retry_after:
+                    try:
+                        return float(retry_after)
+                    except ValueError:
+                        pass
+        
+        return None
+    
+    def get_stats(self) -> dict:
+        """Get rate limiter statistics."""
+        with self._lock:
+            now = time.time()
+            self._clean_windows(now)
+            
+            return {
+                "requests_per_minute_limit": self.requests_per_minute,
+                "requests_per_second_limit": self.requests_per_second,
+                "current_minute_count": len(self._minute_window),
+                "current_second_count": len(self._second_window),
+                "retry_after_active": (
+                    self._retry_after_until is not None and
+                    now < self._retry_after_until
+                ),
+                "retry_after_remaining": (
+                    max(0, self._retry_after_until - now)
+                    if self._retry_after_until
+                    else 0
+                ),
+            }
+    
+    def reset(self) -> None:
+        """Reset all rate limit state."""
+        with self._lock:
+            self._minute_window.clear()
+            self._second_window.clear()
+            self._retry_after_until = None

agent_tool_resilience/resilient_tool.py

@@ -0,0 +1,393 @@
+"""
+Main decorator and class for resilient tool execution.
+"""
+
+import asyncio
+import functools
+import inspect
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional, Sequence, Type, Union
+
+from .retry import RetryPolicy, RetryError
+from .circuit_breaker import CircuitBreaker, CircuitBreakerOpen
+from .fallback import FallbackChain, FallbackError
+from .validator import ResultValidator, ValidationError
+from .tracer import ToolExecutionTracer
+from .rate_limit import RateLimitHandler
+
+
+@dataclass
+class ResilientTool:
+    """
+    Decorator/wrapper for making tool calls resilient.
+    
+    Combines retry, circuit breaker, fallback, validation, and rate limiting
+    into a single easy-to-use interface.
+    
+    Attributes:
+        retry: Retry policy for transient failures
+        circuit_breaker: Circuit breaker for cascade prevention
+        fallback: Fallback function or chain for graceful degradation
+        validator: Result validator for output verification
+        rate_limit: Rate limit handler for API throttling
+        tracer: Execution tracer for observability
+        name: Name for this tool (used in tracing)
+        timeout: Timeout in seconds (None = no timeout)
+        on_error: Callback for errors (before retry/fallback)
+    """
+    retry: Optional[RetryPolicy] = None
+    circuit_breaker: Optional[CircuitBreaker] = None
+    fallback: Optional[Union[Callable[..., Any], FallbackChain]] = None
+    validator: Optional[ResultValidator] = None
+    rate_limit: Optional[RateLimitHandler] = None
+    tracer: Optional[ToolExecutionTracer] = None
+    name: Optional[str] = None
+    timeout: Optional[float] = None
+    on_error: Optional[Callable[[Exception], None]] = None
+    
+    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
+        """Use as decorator."""
+        tool_name = self.name or func.__name__
+        
+        if asyncio.iscoroutinefunction(func):
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                return await self._execute_async(func, tool_name, *args, **kwargs)
+            return async_wrapper
+        else:
+            @functools.wraps(func)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                return self._execute(func, tool_name, *args, **kwargs)
+            return sync_wrapper
+    
+    def _execute(
+        self,
+        func: Callable[..., Any],
+        tool_name: str,
+        *args: Any,
+        **kwargs: Any
+    ) -> Any:
+        """Execute the function with all resilience features."""
+        attempt = 0
+        last_exception: Optional[Exception] = None
+        exec_id: Optional[str] = None
+        
+        # Check circuit breaker first
+        if self.circuit_breaker:
+            if self.circuit_breaker.is_open:
+                if self.tracer:
+                    self.tracer.record_circuit_open(tool_name)
+                
+                # Try fallback
+                if self.fallback:
+                    return self._execute_fallback(
+                        tool_name, args, kwargs,
+                        CircuitBreakerOpen("Circuit is open")
+                    )
+                
+                raise CircuitBreakerOpen("Circuit breaker is open")
+        
+        max_attempts = self.retry.max_attempts if self.retry else 1
+        
+        while attempt < max_attempts:
+            attempt += 1
+            
+            try:
+                # Rate limiting
+                if self.rate_limit:
+                    self.rate_limit.acquire()
+                
+                # Start tracing
+                if self.tracer:
+                    exec_id = self.tracer.start_execution(tool_name, attempt)
+                
+                # Execute with timeout
+                if self.timeout:
+                    import signal
+                    
+                    def timeout_handler(signum, frame):
+                        raise TimeoutError(f"Tool {tool_name} timed out after {self.timeout}s")
+                    
+                    old_handler = signal.signal(signal.SIGALRM, timeout_handler)
+                    signal.alarm(int(self.timeout))
+                    try:
+                        result = func(*args, **kwargs)
+                    finally:
+                        signal.alarm(0)
+                        signal.signal(signal.SIGALRM, old_handler)
+                else:
+                    result = func(*args, **kwargs)
+                
+                # Validate result
+                if self.validator:
+                    result = self.validator.validate(result)
+                
+                # Record success
+                if self.circuit_breaker:
+                    self.circuit_breaker._record_success()
+                
+                if self.tracer:
+                    self.tracer.record_success(exec_id, tool_name, attempt, result)
+                
+                return result
+                
+            except Exception as e:
+                last_exception = e
+                
+                if self.on_error:
+                    self.on_error(e)
+                
+                # Record failure in circuit breaker
+                if self.circuit_breaker:
+                    self.circuit_breaker._record_failure(e)
+                
+                # Check if we should retry
+                should_retry = (
+                    self.retry and
+                    attempt < max_attempts and
+                    self.retry.should_retry(e)
+                )
+                
+                # Record in tracer
+                if self.tracer:
+                    self.tracer.record_failure(
+                        exec_id, tool_name, attempt, e,
+                        will_retry=should_retry
+                    )
+                
+                if should_retry:
+                    # Calculate delay
+                    delay = self.retry.calculate_delay(attempt - 1)
+                    
+                    if self.retry.on_retry:
+                        self.retry.on_retry(attempt, e, delay)
+                    
+                    import time
+                    time.sleep(delay)
+                else:
+                    break
+        
+        # All retries exhausted, try fallback
+        if self.fallback:
+            return self._execute_fallback(tool_name, args, kwargs, last_exception)
+        
+        # No fallback, raise the error
+        if self.retry and last_exception:
+            raise RetryError(
+                f"All {max_attempts} attempts failed for {tool_name}",
+                attempts=max_attempts,
+                last_exception=last_exception
+            )
+        
+        raise last_exception
+    
+    async def _execute_async(
+        self,
+        func: Callable[..., Any],
+        tool_name: str,
+        *args: Any,
+        **kwargs: Any
+    ) -> Any:
+        """Execute an async function with all resilience features."""
+        attempt = 0
+        last_exception: Optional[Exception] = None
+        exec_id: Optional[str] = None
+        
+        # Check circuit breaker first
+        if self.circuit_breaker:
+            if self.circuit_breaker.is_open:
+                if self.tracer:
+                    self.tracer.record_circuit_open(tool_name)
+                
+                # Try fallback
+                if self.fallback:
+                    return await self._execute_fallback_async(
+                        tool_name, args, kwargs,
+                        CircuitBreakerOpen("Circuit is open")
+                    )
+                
+                raise CircuitBreakerOpen("Circuit breaker is open")
+        
+        max_attempts = self.retry.max_attempts if self.retry else 1
+        
+        while attempt < max_attempts:
+            attempt += 1
+            
+            try:
+                # Rate limiting
+                if self.rate_limit:
+                    await self.rate_limit.acquire_async()
+                
+                # Start tracing
+                if self.tracer:
+                    exec_id = self.tracer.start_execution(tool_name, attempt)
+                
+                # Execute with timeout
+                if self.timeout:
+                    result = await asyncio.wait_for(
+                        func(*args, **kwargs),
+                        timeout=self.timeout
+                    )
+                else:
+                    result = await func(*args, **kwargs)
+                
+                # Validate result
+                if self.validator:
+                    result = self.validator.validate(result)
+                
+                # Record success
+                if self.circuit_breaker:
+                    self.circuit_breaker._record_success()
+                
+                if self.tracer:
+                    self.tracer.record_success(exec_id, tool_name, attempt, result)
+                
+                return result
+                
+            except Exception as e:
+                last_exception = e
+                
+                if self.on_error:
+                    self.on_error(e)
+                
+                # Record failure in circuit breaker
+                if self.circuit_breaker:
+                    self.circuit_breaker._record_failure(e)
+                
+                # Check if we should retry
+                should_retry = (
+                    self.retry and
+                    attempt < max_attempts and
+                    self.retry.should_retry(e)
+                )
+                
+                # Record in tracer
+                if self.tracer:
+                    self.tracer.record_failure(
+                        exec_id, tool_name, attempt, e,
+                        will_retry=should_retry
+                    )
+                
+                if should_retry:
+                    # Calculate delay
+                    delay = self.retry.calculate_delay(attempt - 1)
+                    
+                    if self.retry.on_retry:
+                        self.retry.on_retry(attempt, e, delay)
+                    
+                    await asyncio.sleep(delay)
+                else:
+                    break
+        
+        # All retries exhausted, try fallback
+        if self.fallback:
+            return await self._execute_fallback_async(
+                tool_name, args, kwargs, last_exception
+            )
+        
+        # No fallback, raise the error
+        if self.retry and last_exception:
+            raise RetryError(
+                f"All {max_attempts} attempts failed for {tool_name}",
+                attempts=max_attempts,
+                last_exception=last_exception
+            )
+        
+        raise last_exception
+    
+    def _execute_fallback(
+        self,
+        tool_name: str,
+        args: tuple,
+        kwargs: dict,
+        primary_exception: Exception
+    ) -> Any:
+        """Execute the fallback."""
+        if isinstance(self.fallback, FallbackChain):
+            result = self.fallback.execute(
+                *args,
+                primary_exception=primary_exception,
+                **kwargs
+            )
+        else:
+            result = self.fallback(*args, **kwargs)
+        
+        if self.tracer:
+            fallback_name = (
+                self.fallback.__name__
+                if callable(self.fallback) and hasattr(self.fallback, "__name__")
+                else "fallback"
+            )
+            self.tracer.record_fallback(tool_name, fallback_name, result)
+        
+        return result
+    
+    async def _execute_fallback_async(
+        self,
+        tool_name: str,
+        args: tuple,
+        kwargs: dict,
+        primary_exception: Exception
+    ) -> Any:
+        """Execute the fallback asynchronously."""
+        if isinstance(self.fallback, FallbackChain):
+            if asyncio.iscoroutinefunction(self.fallback.fallbacks[0]):
+                result = await self.fallback.execute_async(
+                    *args,
+                    primary_exception=primary_exception,
+                    **kwargs
+                )
+            else:
+                result = self.fallback.execute(
+                    *args,
+                    primary_exception=primary_exception,
+                    **kwargs
+                )
+        else:
+            if asyncio.iscoroutinefunction(self.fallback):
+                result = await self.fallback(*args, **kwargs)
+            else:
+                result = self.fallback(*args, **kwargs)
+        
+        if self.tracer:
+            fallback_name = (
+                self.fallback.__name__
+                if callable(self.fallback) and hasattr(self.fallback, "__name__")
+                else "fallback"
+            )
+            self.tracer.record_fallback(tool_name, fallback_name, result)
+        
+        return result
+
+
+def resilient_tool(
+    retry: Optional[RetryPolicy] = None,
+    circuit_breaker: Optional[CircuitBreaker] = None,
+    fallback: Optional[Union[Callable[..., Any], FallbackChain]] = None,
+    validator: Optional[ResultValidator] = None,
+    rate_limit: Optional[RateLimitHandler] = None,
+    tracer: Optional[ToolExecutionTracer] = None,
+    name: Optional[str] = None,
+    timeout: Optional[float] = None,
+    on_error: Optional[Callable[[Exception], None]] = None
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """
+    Decorator factory for resilient tools.
+    
+    Usage:
+        @resilient_tool(retry=RetryPolicy(max_attempts=3))
+        def my_tool():
+            ...
+    """
+    wrapper = ResilientTool(
+        retry=retry,
+        circuit_breaker=circuit_breaker,
+        fallback=fallback,
+        validator=validator,
+        rate_limit=rate_limit,
+        tracer=tracer,
+        name=name,
+        timeout=timeout,
+        on_error=on_error
+    )
+    return wrapper.__call__