chuk-tool-processor 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

chuk_tool_processor/execution/wrappers/rate_limiting.py

@@ -1,149 +1,262 @@
 # chuk_tool_processor/execution/wrappers/rate_limiting.py
+"""
+Async-native rate-limiting wrapper.
+
+Two layers of limits are enforced:
+
+* **Global** - ``<N requests> / <period>`` over *all* tools.
+* **Per-tool** - independent ``<N requests> / <period>`` windows.
+
+A simple sliding-window algorithm with timestamp queues is used.  
+`asyncio.Lock` guards shared state so the wrapper can be used safely from
+multiple coroutines.
+"""
+from __future__ import annotations
+
 import asyncio
+import inspect
 import time
-from datetime import datetime
-from typing import Dict, Optional, List, Any, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
-# imports
 from chuk_tool_processor.models.tool_call import ToolCall
 from chuk_tool_processor.models.tool_result import ToolResult
-from chuk_tool_processor.core.exceptions import ToolExecutionError
+from chuk_tool_processor.logging import get_logger
 
+logger = get_logger("chuk_tool_processor.execution.wrappers.rate_limiting")
 
+# --------------------------------------------------------------------------- #
+# Core limiter
+# --------------------------------------------------------------------------- #
 class RateLimiter:
     """
-    Rate limiter for tool executions.
-    Supports per-tool rate limits and global rate limits.
+    Async-native rate limiter for controlling execution frequency.
+    
+    Implements a sliding window algorithm to enforce rate limits both globally
+    and per-tool. All operations are thread-safe using asyncio locks.
     """
+    
     def __init__(
         self,
+        *,
         global_limit: Optional[int] = None,
         global_period: float = 60.0,
-        tool_limits: Optional[Dict[str, Tuple[int, float]]] = None
-    ):
+        tool_limits: Optional[Dict[str, Tuple[int, float]]] = None,
+    ) -> None:
         """
         Initialize the rate limiter.
+        
+        Args:
+            global_limit: Maximum global requests per period (None = no limit)
+            global_period: Time period in seconds for the global limit
+            tool_limits: Dict mapping tool names to (limit, period) tuples
         """
         self.global_limit = global_limit
         self.global_period = global_period
         self.tool_limits = tool_limits or {}
-        
-        # Track request timestamps
-        self._global_timestamps: List[float] = []
-        self._tool_timestamps: Dict[str, List[float]] = {}
-        
-        # Locks for concurrency safety
+
+        # Timestamp queues
+        self._global_ts: List[float] = []
+        self._tool_ts: Dict[str, List[float]] = {}
+
+        # Locks for thread safety
         self._global_lock = asyncio.Lock()
         self._tool_locks: Dict[str, asyncio.Lock] = {}
-    
-    async def _wait_for_global_limit(self) -> None:
-        """
-        Wait until global rate limit allows another request.
-        """
+        
+        logger.debug(
+            f"Initialized rate limiter: global={global_limit}/{global_period}s, "
+            f"tool-specific={len(self.tool_limits)} tools"
+        )
+
+    # --------------------- helpers -------------------- #
+    async def _acquire_global(self) -> None:
+        """Block until a global slot is available."""
         if self.global_limit is None:
             return
-        
+
         while True:
-            # Acquire lock to check and possibly record
             async with self._global_lock:
-                now = time.time()
-                # Remove expired timestamps
+                now = time.monotonic()
                 cutoff = now - self.global_period
-                self._global_timestamps = [ts for ts in self._global_timestamps if ts > cutoff]
-                # If under limit, record and proceed
-                if len(self._global_timestamps) < self.global_limit:
-                    self._global_timestamps.append(now)
+                
+                # Prune expired timestamps
+                self._global_ts = [t for t in self._global_ts if t > cutoff]
+
+                # Check if we're under the limit
+                if len(self._global_ts) < self.global_limit:
+                    self._global_ts.append(now)
                     return
-                # Otherwise compute wait time
-                oldest = min(self._global_timestamps)
-                wait_time = (oldest + self.global_period) - now
-            # Sleep outside lock
-            if wait_time > 0:
-                await asyncio.sleep(wait_time)
-            else:
-                # retry immediately
-                continue
-    
-    async def _wait_for_tool_limit(self, tool: str) -> None:
-        """
-        Wait until tool-specific rate limit allows another request.
-        """
-        # Check if tool has a limit
+
+                # Calculate wait time until a slot becomes available
+                wait = (self._global_ts[0] + self.global_period) - now
+                
+            logger.debug(f"Global rate limit reached, waiting {wait:.2f}s")
+            await asyncio.sleep(wait)
+
+    async def _acquire_tool(self, tool: str) -> None:
+        """Block until a per-tool slot is available (if the tool has a limit)."""
         if tool not in self.tool_limits:
             return
+
         limit, period = self.tool_limits[tool]
-        
-        # Initialize lock and timestamps list if needed
-        if tool not in self._tool_locks:
-            self._tool_locks[tool] = asyncio.Lock()
-        if tool not in self._tool_timestamps:
-            self._tool_timestamps[tool] = []
-        
+        lock = self._tool_locks.setdefault(tool, asyncio.Lock())
+        buf = self._tool_ts.setdefault(tool, [])
+
         while True:
-            async with self._tool_locks[tool]:
-                now = time.time()
-                # Remove expired timestamps
+            async with lock:
+                now = time.monotonic()
                 cutoff = now - period
-                self._tool_timestamps[tool] = [ts for ts in self._tool_timestamps[tool] if ts > cutoff]
-                # If under limit, record and proceed
-                if len(self._tool_timestamps[tool]) < limit:
-                    self._tool_timestamps[tool].append(now)
+                
+                # Prune expired timestamps in-place
+                buf[:] = [t for t in buf if t > cutoff]
+
+                # Check if we're under the limit
+                if len(buf) < limit:
+                    buf.append(now)
                     return
-                # Otherwise compute wait time
-                oldest = min(self._tool_timestamps[tool])
-                wait_time = (oldest + period) - now
-            # Sleep outside lock
-            if wait_time > 0:
-                await asyncio.sleep(wait_time)
-            else:
-                continue
-    
+
+                # Calculate wait time until a slot becomes available
+                wait = (buf[0] + period) - now
+                
+            logger.debug(f"Tool '{tool}' rate limit reached, waiting {wait:.2f}s")
+            await asyncio.sleep(wait)
+
+    # ----------------------- public -------------------- #
     async def wait(self, tool: str) -> None:
         """
-        Wait until rate limits allow execution of the given tool.
+        Block until rate limits allow execution.
+        
+        This method blocks until both global and tool-specific rate limits
+        allow one more execution of the specified tool.
+        
+        Args:
+            tool: Name of the tool being executed
+        """
+        await self._acquire_global()
+        await self._acquire_tool(tool)
+        
+    async def check_limits(self, tool: str) -> Tuple[bool, bool]:
+        """
+        Check if the tool would be rate limited without consuming a slot.
+        
+        This is a non-blocking method useful for checking limits without
+        affecting the rate limiting state.
+        
+        Args:
+            tool: Name of the tool to check
+            
+        Returns:
+            Tuple of (global_limit_reached, tool_limit_reached)
         """
-        # Wait for global limit first
-        await self._wait_for_global_limit()
-        # Then wait for tool-specific limit
-        await self._wait_for_tool_limit(tool)
+        global_limited = False
+        tool_limited = False
+        
+        # Check global limit
+        if self.global_limit is not None:
+            async with self._global_lock:
+                now = time.monotonic()
+                cutoff = now - self.global_period
+                active_ts = [t for t in self._global_ts if t > cutoff]
+                global_limited = len(active_ts) >= self.global_limit
+        
+        # Check tool limit
+        if tool in self.tool_limits:
+            limit, period = self.tool_limits[tool]
+            async with self._tool_locks.setdefault(tool, asyncio.Lock()):
+                now = time.monotonic()
+                cutoff = now - period
+                buf = self._tool_ts.get(tool, [])
+                active_ts = [t for t in buf if t > cutoff]
+                tool_limited = len(active_ts) >= limit
+                
+        return global_limited, tool_limited
 
 
+# --------------------------------------------------------------------------- #
+# Executor wrapper
+# --------------------------------------------------------------------------- #
 class RateLimitedToolExecutor:
     """
-    Wrapper for a tool executor that applies rate limiting.
+    Executor wrapper that applies rate limiting to tool executions.
+    
+    This wrapper delegates to another executor but ensures that all
+    tool calls respect the configured rate limits.
     """
-    def __init__(
-        self,
-        executor: Any,
-        rate_limiter: RateLimiter
-    ):
+
+    def __init__(self, executor: Any, limiter: RateLimiter) -> None:
         """
         Initialize the rate-limited executor.
+        
+        Args:
+            executor: The underlying executor to wrap
+            limiter: The RateLimiter that controls execution frequency
         """
         self.executor = executor
-        self.rate_limiter = rate_limiter
-    
+        self.limiter = limiter
+        logger.debug(f"Initialized rate-limited executor")
+
     async def execute(
         self,
         calls: List[ToolCall],
-        timeout: Optional[float] = None
+        timeout: Optional[float] = None,
+        use_cache: bool = True,
     ) -> List[ToolResult]:
         """
-        Execute tool calls with rate limiting.
+        Execute tool calls while respecting rate limits.
+        
+        This method blocks until rate limits allow execution, then delegates
+        to the underlying executor.
+        
+        Args:
+            calls: List of tool calls to execute
+            timeout: Optional timeout for execution
+            use_cache: Whether to use cached results (forwarded to underlying executor)
+            
+        Returns:
+            List of tool results
         """
-        # Apply rate limiting to each call
-        for call in calls:
-            await self.rate_limiter.wait(call.tool)
-        # Delegate to inner executor
+        if not calls:
+            return []
+            
+        # Block for each call *before* dispatching to the wrapped executor
+        for c in calls:
+            await self.limiter.wait(c.tool)
+            
+        # Check if the executor has a use_cache parameter
+        if hasattr(self.executor, "execute"):
+            sig = inspect.signature(self.executor.execute)
+            if "use_cache" in sig.parameters:
+                return await self.executor.execute(calls, timeout=timeout, use_cache=use_cache)
+                
+        # Fall back to standard execute method
         return await self.executor.execute(calls, timeout=timeout)
 
 
+# --------------------------------------------------------------------------- #
+# Convenience decorator for tools
+# --------------------------------------------------------------------------- #
 def rate_limited(limit: int, period: float = 60.0):
     """
-    Decorator to specify rate limits for a tool class.
+    Class decorator that marks a Tool with default rate-limit metadata.
+    
+    This allows higher-level code to detect and configure rate limiting
+    for the tool class.
+
+    Example:
+        @rate_limited(limit=10, period=60.0)
+        class WeatherTool:
+            async def execute(self, location: str) -> Dict[str, Any]:
+                # Implementation
+                
+    Args:
+        limit: Maximum number of calls allowed in the period
+        period: Time period in seconds
+        
+    Returns:
+        Decorated class with rate limit metadata
     """
     def decorator(cls):
         cls._rate_limit = limit
         cls._rate_period = period
         return cls
-    return decorator
+
+    return decorator

chuk_tool_processor/execution/wrappers/retry.py

@@ -1,20 +1,36 @@
-# chuk_tool_processor/retry.py
+# chuk_tool_processor/execution/wrappers/retry.py
+"""
+Async-native retry wrapper for tool execution.
+
+This module provides a retry mechanism for tool calls that can automatically
+retry failed executions based on configurable criteria and backoff strategies.
+"""
+from __future__ import annotations
+
 import asyncio
 import logging
 import random
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type, Union
 
-# imports
 from chuk_tool_processor.models.tool_call import ToolCall
 from chuk_tool_processor.models.tool_result import ToolResult
+from chuk_tool_processor.logging import get_logger
 
-logger = logging.getLogger(__name__)
+logger = get_logger("chuk_tool_processor.execution.wrappers.retry")
 
 
 class RetryConfig:
     """
     Configuration for retry behavior.
+    
+    Attributes:
+        max_retries: Maximum number of retry attempts
+        base_delay: Base delay between retries in seconds
+        max_delay: Maximum delay between retries in seconds
+        jitter: Whether to add random jitter to delays
+        retry_on_exceptions: List of exception types to retry on
+        retry_on_error_substrings: List of error message substrings to retry on
     """
     def __init__(
         self,
@@ -33,6 +49,17 @@ class RetryConfig:
         self.retry_on_error_substrings = retry_on_error_substrings or []
     
     def should_retry(self, attempt: int, error: Optional[Exception] = None, error_str: Optional[str] = None) -> bool:
+        """
+        Determine if a retry should be attempted.
+        
+        Args:
+            attempt: Current attempt number (0-based)
+            error: Exception that caused the failure, if any
+            error_str: Error message string, if any
+            
+        Returns:
+            True if a retry should be attempted, False otherwise
+        """
         if attempt >= self.max_retries:
             return False
         if not self.retry_on_exceptions and not self.retry_on_error_substrings:
@@ -44,6 +71,15 @@ class RetryConfig:
         return False
     
     def get_delay(self, attempt: int) -> float:
+        """
+        Calculate the delay for the current attempt with exponential backoff.
+        
+        Args:
+            attempt: Current attempt number (0-based)
+            
+        Returns:
+            Delay in seconds
+        """
         delay = min(self.base_delay * (2 ** attempt), self.max_delay)
         if self.jitter:
             delay *= (0.5 + random.random())
@@ -53,29 +89,58 @@ class RetryConfig:
 class RetryableToolExecutor:
     """
     Wrapper for a tool executor that applies retry logic.
+    
+    This executor wraps another executor and automatically retries failed
+    tool calls based on configured retry policies.
     """
     def __init__(
         self,
         executor: Any,
-        default_config: RetryConfig = None,
-        tool_configs: Dict[str, RetryConfig] = None
+        default_config: Optional[RetryConfig] = None,
+        tool_configs: Optional[Dict[str, RetryConfig]] = None
     ):
+        """
+        Initialize the retryable executor.
+        
+        Args:
+            executor: The underlying executor to wrap
+            default_config: Default retry configuration for all tools
+            tool_configs: Tool-specific retry configurations
+        """
         self.executor = executor
         self.default_config = default_config or RetryConfig()
         self.tool_configs = tool_configs or {}
     
     def _get_config(self, tool: str) -> RetryConfig:
+        """Get the retry configuration for a specific tool."""
         return self.tool_configs.get(tool, self.default_config)
     
     async def execute(
         self,
         calls: List[ToolCall],
-        timeout: Optional[float] = None
+        timeout: Optional[float] = None,
+        use_cache: bool = True
     ) -> List[ToolResult]:
+        """
+        Execute tool calls with retry logic.
+        
+        Args:
+            calls: List of tool calls to execute
+            timeout: Optional timeout for each execution
+            use_cache: Whether to use cached results (passed to underlying executor)
+            
+        Returns:
+            List of tool results
+        """
+        # Handle empty calls list
+        if not calls:
+            return []
+            
+        # Execute each call with retries
         results: List[ToolResult] = []
         for call in calls:
             config = self._get_config(call.tool)
-            result = await self._execute_with_retry(call, config, timeout)
+            result = await self._execute_with_retry(call, config, timeout, use_cache)
             results.append(result)
         return results
     
@@ -83,8 +148,21 @@ class RetryableToolExecutor:
         self,
         call: ToolCall,
         config: RetryConfig,
-        timeout: Optional[float]
+        timeout: Optional[float],
+        use_cache: bool
     ) -> ToolResult:
+        """
+        Execute a single tool call with retries.
+        
+        Args:
+            call: Tool call to execute
+            config: Retry configuration to use
+            timeout: Optional timeout for execution
+            use_cache: Whether to use cached results
+            
+        Returns:
+            Tool result after retries
+        """
         attempt = 0
         last_error: Optional[str] = None
         pid = 0
@@ -92,24 +170,31 @@ class RetryableToolExecutor:
         
         while True:
             start_time = datetime.now(timezone.utc)
+            
             try:
-                # execute call
-                tool_results = await self.executor.execute([call], timeout=timeout)
+                # Pass the use_cache parameter if the executor supports it
+                executor_kwargs = {"timeout": timeout}
+                if hasattr(self.executor, "use_cache"):
+                    executor_kwargs["use_cache"] = use_cache
+                
+                # Execute call
+                tool_results = await self.executor.execute([call], **executor_kwargs)
                 result = tool_results[0]
                 pid = result.pid
                 machine = result.machine
                 
-                # error in result
+                # Check for error in result
                 if result.error:
                     last_error = result.error
                     if config.should_retry(attempt, error_str=result.error):
                         logger.debug(
-                            f"Retrying tool {call.tool} after error: {result.error} (attempt {attempt + 1})"
+                            f"Retrying tool {call.tool} after error: {result.error} (attempt {attempt + 1}/{config.max_retries})"
                         )
                         await asyncio.sleep(config.get_delay(attempt))
                         attempt += 1
                         continue
-                    # no retry: if any retries happened, wrap final error
+                        
+                    # No retry: if any retries happened, wrap final error
                     if attempt > 0:
                         end_time = datetime.now(timezone.utc)
                         final = ToolResult(
@@ -121,26 +206,31 @@ class RetryableToolExecutor:
                             machine=machine,
                             pid=pid
                         )
-                        # attach attempts
-                        object.__setattr__(final, 'attempts', attempt)
+                        # Attach attempts
+                        final.attempts = attempt + 1  # Include the original attempt
                         return final
-                    # no retries occurred, return the original failure
+                        
+                    # No retries occurred, return the original failure
+                    result.attempts = 1
                     return result
                 
-                # success: attach attempts and return
-                object.__setattr__(result, 'attempts', attempt)
+                # Success: attach attempts and return
+                result.attempts = attempt + 1  # Include the original attempt
                 return result
+                
             except Exception as e:
                 err_str = str(e)
                 last_error = err_str
+                
                 if config.should_retry(attempt, error=e):
                     logger.info(
-                        f"Retrying tool {call.tool} after exception: {err_str} (attempt {attempt + 1})"
+                        f"Retrying tool {call.tool} after exception: {err_str} (attempt {attempt + 1}/{config.max_retries})"
                     )
                     await asyncio.sleep(config.get_delay(attempt))
                     attempt += 1
                     continue
-                # no more retries: return error result
+                    
+                # No more retries: return error result
                 end_time = datetime.now(timezone.utc)
                 final_exc = ToolResult(
                     tool=call.tool,
@@ -151,7 +241,7 @@ class RetryableToolExecutor:
                     machine=machine,
                     pid=pid
                 )
-                object.__setattr__(final_exc, 'attempts', attempt + 1)
+                final_exc.attempts = attempt + 1  # Include the original attempt
                 return final_exc
 
 
@@ -163,6 +253,26 @@ def retryable(
     retry_on_exceptions: Optional[List[Type[Exception]]] = None,
     retry_on_error_substrings: Optional[List[str]] = None
 ):
+    """
+    Decorator for tool classes to configure retry behavior.
+    
+    Example:
+        @retryable(max_retries=5, base_delay=2.0)
+        class MyTool:
+            async def execute(self, x: int, y: int) -> int:
+                return x + y
+                
+    Args:
+        max_retries: Maximum number of retry attempts
+        base_delay: Base delay between retries in seconds
+        max_delay: Maximum delay between retries in seconds
+        jitter: Whether to add random jitter to delays
+        retry_on_exceptions: List of exception types to retry on
+        retry_on_error_substrings: List of error message substrings to retry on
+        
+    Returns:
+        Decorated class with retry configuration
+    """
     def decorator(cls):
         cls._retry_config = RetryConfig(
             max_retries=max_retries,
@@ -173,4 +283,4 @@ def retryable(
             retry_on_error_substrings=retry_on_error_substrings
         )
         return cls
-    return decorator
+    return decorator