edsl 0.1.46__py3-none-any.whl → 0.1.48__py3-none-any.whl

edsl/buckets/token_bucket_client.py

@@ -0,0 +1,509 @@
+"""
+Token Bucket Client for distributed rate limiting in EDSL.
+
+This module provides a client implementation for interacting with a remote token
+bucket server. It implements the same interface as TokenBucket, but delegates
+operations to a remote server, enabling distributed rate limiting across
+multiple processes or machines.
+"""
+
+from typing import Union, Optional, Dict, List, Any, Tuple
+import asyncio
+import time
+import aiohttp
+from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
+
+
+class TokenBucketClient:
+    """
+    Client implementation for interacting with a remote token bucket server.
+    
+    TokenBucketClient implements the same interface as TokenBucket, but
+    delegates operations to a remote server via REST API calls. This enables
+    distributed rate limiting across multiple processes or machines, ensuring
+    that rate limits are properly enforced in distributed environments.
+    
+    The client maintains minimal local state and fetches most information from
+    the server when needed. It creates the bucket on the server during
+    initialization if it doesn't already exist.
+    
+    Attributes:
+        bucket_name (str): Name identifier for the bucket (usually service name)
+        bucket_type (str): Type of bucket ("requests" or "tokens")
+        capacity (float): Maximum tokens the bucket can hold
+        refill_rate (float): Rate at which tokens are refilled (tokens per second)
+        api_base_url (str): Base URL for the token bucket server API
+        bucket_id (str): Unique identifier for this bucket on the server
+        creation_time (float): Local timestamp when this client was created
+        turbo_mode (bool): Flag indicating if turbo mode is active
+        
+    Example:
+        >>> # Create a client connected to a running token bucket server
+        >>> client = TokenBucketClient(
+        ...     bucket_name="openai", 
+        ...     bucket_type="requests",
+        ...     capacity=100, 
+        ...     refill_rate=10,
+        ...     api_base_url="http://localhost:8000"
+        ... )
+        >>> # Now use this client just like a regular TokenBucket
+    """
+
+    def __init__(
+        self,
+        *,
+        bucket_name: str,
+        bucket_type: str,
+        capacity: Union[int, float],
+        refill_rate: Union[int, float],
+        api_base_url: str = "http://localhost:8000",
+    ):
+        """
+        Initialize a new TokenBucketClient connected to a remote token bucket server.
+        
+        Creates a new TokenBucketClient instance that connects to a remote token
+        bucket server. During initialization, it attempts to create the bucket
+        on the server if it doesn't already exist.
+        
+        Args:
+            bucket_name: Name identifier for the bucket (usually service name)
+            bucket_type: Type of bucket, either "requests" or "tokens" 
+            capacity: Maximum tokens the bucket can hold
+            refill_rate: Rate at which tokens are added (tokens per second)
+            api_base_url: Base URL for the token bucket server API
+                         (default: "http://localhost:8000")
+                         
+        Raises:
+            ValueError: If bucket creation on the server fails
+            
+        Example:
+            >>> client = TokenBucketClient(
+            ...     bucket_name="openai",
+            ...     bucket_type="requests", 
+            ...     capacity=100,
+            ...     refill_rate=10
+            ... )
+        """
+        self.bucket_name = bucket_name
+        self.bucket_type = bucket_type
+        self.capacity = capacity
+        self.refill_rate = refill_rate
+        self.api_base_url = api_base_url
+        self.bucket_id = f"{bucket_name}_{bucket_type}"
+
+        # Initialize the bucket on the server
+        asyncio.run(self._create_bucket())
+
+        # Cache some values locally
+        self.creation_time = time.monotonic()
+        self.turbo_mode = False
+
+    async def _create_bucket(self) -> None:
+        """
+        Create or retrieve the bucket on the remote server.
+        
+        This private async method sends a request to the server to create a new
+        bucket with the specified parameters. If the bucket already exists on
+        the server, it updates the local parameters to match the server's values.
+        
+        Raises:
+            ValueError: If the server returns an error
+        """
+        async with aiohttp.ClientSession() as session:
+            # Prepare payload with bucket parameters
+            payload = {
+                "bucket_name": self.bucket_name,
+                "bucket_type": self.bucket_type,
+                "capacity": self.capacity,
+                "refill_rate": self.refill_rate,
+            }
+            
+            # Send request to create/retrieve bucket
+            async with session.post(
+                f"{self.api_base_url}/bucket",
+                json=payload,
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(f"Unexpected error: {await response.text()}")
+
+                # Process server response
+                result = await response.json()
+                if result["status"] == "existing":
+                    # Update our local values to match the existing bucket
+                    self.capacity = float(result["bucket"]["capacity"])
+                    self.refill_rate = float(result["bucket"]["refill_rate"])
+
+    def turbo_mode_on(self) -> None:
+        """
+        Enable turbo mode to bypass rate limits.
+        
+        Turbo mode sets the refill rate to infinity on the server,
+        effectively bypassing rate limits. This is useful for testing
+        or when rate limits are not needed.
+        
+        Raises:
+            ValueError: If the server returns an error
+            
+        Example:
+            >>> client = TokenBucketClient(bucket_name="test", bucket_type="test", 
+            ...                           capacity=100, refill_rate=10)
+            >>> client.turbo_mode_on()  # Now rate limits are bypassed
+        """
+        asyncio.run(self._set_turbo_mode(True))
+        self.turbo_mode = True
+
+    def turbo_mode_off(self) -> None:
+        """
+        Disable turbo mode and restore original rate limits.
+        
+        This method restores the original refill rates on the server,
+        re-enabling rate limiting after it was bypassed with turbo_mode_on().
+        
+        Raises:
+            ValueError: If the server returns an error
+            
+        Example:
+            >>> client = TokenBucketClient(bucket_name="test", bucket_type="test", 
+            ...                           capacity=100, refill_rate=10)
+            >>> client.turbo_mode_on()  # Bypass rate limits
+            >>> # Do some work without rate limiting
+            >>> client.turbo_mode_off()  # Restore rate limits
+        """
+        asyncio.run(self._set_turbo_mode(False))
+        self.turbo_mode = False
+
+    async def add_tokens(self, amount: Union[int, float]) -> None:
+        """
+        Add tokens to the bucket on the server.
+        
+        This async method adds tokens to the bucket on the server.
+        It's useful for manually restoring tokens or increasing the 
+        available tokens beyond the normal refill rate.
+        
+        Args:
+            amount: Number of tokens to add to the bucket
+            
+        Raises:
+            ValueError: If the server returns an error
+            
+        Example:
+            >>> import asyncio
+            >>> client = TokenBucketClient(bucket_name="test", bucket_type="test", 
+            ...                           capacity=100, refill_rate=10)
+            >>> # Add 50 tokens to the bucket
+            >>> asyncio.run(client.add_tokens(50))
+        """
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.api_base_url}/bucket/{self.bucket_id}/add_tokens",
+                params={"amount": amount},
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(f"Failed to add tokens: {await response.text()}")
+
+    async def _set_turbo_mode(self, state: bool) -> None:
+        """
+        Set the turbo mode state on the server.
+        
+        This private async method sends a request to the server to
+        enable or disable turbo mode.
+        
+        Args:
+            state: True to enable turbo mode, False to disable
+            
+        Raises:
+            ValueError: If the server returns an error
+        """
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.api_base_url}/bucket/{self.bucket_id}/turbo_mode/{str(state).lower()}"
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(
+                        f"Failed to set turbo mode: {await response.text()}"
+                    )
+
+    async def get_tokens(
+        self, amount: Union[int, float] = 1, cheat_bucket_capacity: bool = True
+    ) -> None:
+        """
+        Request tokens from the token bucket on the server.
+        
+        This async method requests tokens from the token bucket on the server.
+        It will either return immediately if tokens are available or raise an
+        exception if tokens are not available.
+        
+        Args:
+            amount: Number of tokens to request (default: 1)
+            cheat_bucket_capacity: If True, allow exceeding capacity temporarily
+                                  (default: True)
+            
+        Raises:
+            ValueError: If the server returns an error, which may indicate
+                      insufficient tokens are available
+                      
+        Example:
+            >>> import asyncio
+            >>> client = TokenBucketClient(bucket_name="test", bucket_type="test", 
+            ...                           capacity=100, refill_rate=10)
+            >>> # Request 20 tokens
+            >>> asyncio.run(client.get_tokens(20))
+        """
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.api_base_url}/bucket/{self.bucket_id}/get_tokens",
+                params={
+                    "amount": amount,
+                    "cheat_bucket_capacity": int(cheat_bucket_capacity),
+                },
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(f"Failed to get tokens: {await response.text()}")
+
+    def get_throughput(self, time_window: Optional[float] = None) -> float:
+        """
+        Calculate the token throughput over a specified time window.
+        
+        This method calculates the average token throughput (tokens per minute)
+        over the specified time window by requesting the bucket status from
+        the server and analyzing token usage.
+        
+        Args:
+            time_window: Time window in seconds to calculate throughput over
+                        (default: entire bucket lifetime)
+                        
+        Returns:
+            Average throughput in tokens per minute
+            
+        Example:
+            >>> client = TokenBucketClient(bucket_name="test", bucket_type="test", 
+            ...                           capacity=100, refill_rate=10)
+            >>> # Calculate throughput over the last 60 seconds
+            >>> throughput = client.get_throughput(60)
+            >>> print(f"Average throughput: {throughput:.1f} tokens/minute")
+        """
+        # Get current bucket status from server
+        status = asyncio.run(self._get_status())
+        now = time.monotonic()
+
+        # Determine start time based on time_window parameter
+        if time_window is None:
+            start_time = self.creation_time
+        else:
+            start_time = now - time_window
+
+        # Ensure start_time isn't before bucket creation
+        if start_time < self.creation_time:
+            start_time = self.creation_time
+
+        # Calculate elapsed time and avoid division by zero
+        elapsed_time = now - start_time
+        if elapsed_time == 0:
+            return status["num_released"] / 0.001  # Avoid division by zero
+
+        # Convert to tokens per minute
+        return (status["num_released"] / elapsed_time) * 60
+
+    async def _get_status(self) -> Dict[str, Any]:
+        """
+        Get the current status of the bucket from the server.
+        
+        This private async method retrieves the current status of the bucket
+        from the server, including the current token count, history log,
+        and various statistics.
+        
+        Returns:
+            Dictionary containing the bucket status information
+            
+        Raises:
+            ValueError: If the server returns an error
+        """
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                f"{self.api_base_url}/bucket/{self.bucket_id}/status"
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(
+                        f"Failed to get bucket status: {await response.text()}"
+                    )
+                return await response.json()
+
+    def __add__(self, other: "TokenBucketClient") -> "TokenBucketClient":
+        """
+        Combine two token bucket clients to create a new one with merged limits.
+        
+        This method combines two token bucket clients by creating a new client
+        with the minimum capacity and refill rate of both inputs. This is useful
+        for creating a client that respects both sets of rate limits.
+        
+        Args:
+            other: Another TokenBucketClient to combine with this one
+            
+        Returns:
+            A new TokenBucketClient with the combined (minimum) limits
+            
+        Example:
+            >>> client1 = TokenBucketClient(bucket_name="service1", bucket_type="requests",
+            ...                            capacity=100, refill_rate=10)
+            >>> client2 = TokenBucketClient(bucket_name="service2", bucket_type="requests",
+            ...                            capacity=50, refill_rate=5)
+            >>> combined = client1 + client2  # Takes the minimum of both limits
+        """
+        return TokenBucketClient(
+            bucket_name=self.bucket_name,
+            bucket_type=self.bucket_type,
+            capacity=min(self.capacity, other.capacity),
+            refill_rate=min(self.refill_rate, other.refill_rate),
+            api_base_url=self.api_base_url,
+        )
+
+    @property
+    def tokens(self) -> float:
+        """
+        Get the current number of tokens available in the bucket.
+        
+        This property retrieves the current token count from the server.
+        
+        Returns:
+            Current number of tokens available in the bucket
+            
+        Example:
+            >>> client = TokenBucketClient(bucket_name="test", bucket_type="test",
+            ...                           capacity=100, refill_rate=10)
+            >>> available = client.tokens
+            >>> print(f"Available tokens: {available}")
+        """
+        status = asyncio.run(self._get_status())
+        return float(status["tokens"])
+
+    def wait_time(self, requested_tokens: Union[float, int]) -> float:
+        """
+        Calculate the time to wait for the requested number of tokens.
+        
+        This method calculates how long to wait (in seconds) for the requested
+        number of tokens to become available, based on the current token count
+        and refill rate.
+        
+        Args:
+            requested_tokens: Number of tokens needed
+            
+        Returns:
+            Time to wait in seconds (0.0 if tokens are already available)
+            
+        Raises:
+            ValueError: If an error occurs while calculating wait time
+            
+        Example:
+            >>> client = TokenBucketClient(bucket_name="test", bucket_type="test",
+            ...                           capacity=100, refill_rate=10)
+            >>> wait_seconds = client.wait_time(50)
+            >>> print(f"Need to wait {wait_seconds:.2f} seconds")
+        """
+        # If we have enough tokens, no need to wait
+        if self.tokens >= float(requested_tokens):
+            return 0.0
+            
+        try:
+            # Calculate time needed to accumulate the required tokens
+            return (requested_tokens - self.tokens) / self.refill_rate
+        except Exception as e:
+            raise ValueError(f"Error calculating wait time: {e}")
+
+    # Note: The commented out method below is a reminder for future implementation
+    # def wait_time(self, num_tokens: Union[int, float]) -> float:
+    #     """Server-side wait time calculation (future implementation)"""
+    #     return 0  # TODO - Need to implement this on the server side
+
+    def visualize(self) -> Figure:
+        """
+        Visualize the token bucket usage over time as a matplotlib figure.
+        
+        This method generates a plot showing the available tokens over time,
+        which can be useful for monitoring and debugging rate limit issues.
+        
+        Returns:
+            A matplotlib Figure object that can be displayed or saved
+            
+        Example:
+            >>> client = TokenBucketClient(bucket_name="test", bucket_type="test",
+            ...                           capacity=100, refill_rate=10)
+            >>> # Do some operations with the bucket
+            >>> plot = client.visualize()
+            >>> # Now you can display or save the plot
+        """
+        # Get the bucket history from the server
+        status = asyncio.run(self._get_status())
+        times, tokens = zip(*status["log"])
+        
+        # Normalize times to start at 0
+        start_time = times[0]
+        times = [t - start_time for t in times]
+
+        # Create the plot
+        fig = plt.figure(figsize=(10, 6))
+        plt.plot(times, tokens, label="Tokens Available")
+        plt.xlabel("Time (seconds)", fontsize=12)
+        plt.ylabel("Number of Tokens", fontsize=12)
+        details = f"{self.bucket_name} ({self.bucket_type}) Bucket Usage Over Time\nCapacity: {self.capacity:.1f}, Refill Rate: {self.refill_rate:.1f}/second"
+        plt.title(details, fontsize=14)
+        plt.legend()
+        plt.grid(True)
+        plt.tight_layout()
+        
+        return fig
+
+
+# Examples and doctests
+if __name__ == "__main__":
+    #import doctest
+
+    # Example showing how to use TokenBucketClient
+    def example_usage():
+        """
+        Example demonstrating how to use TokenBucketClient:
+        
+        ```python
+        import asyncio
+        import time
+        from edsl.buckets.token_bucket_client import TokenBucketClient
+        
+        # Create a client connected to a running token bucket server
+        bucket = TokenBucketClient(
+            bucket_name="openai", 
+            bucket_type="requests",
+            capacity=100, 
+            refill_rate=10,
+            api_base_url="http://localhost:8000"
+        )
+        
+        # Get tokens from the bucket
+        asyncio.run(bucket.get_tokens(50))
+        
+        # Wait for a second
+        time.sleep(1)
+        
+        # Get more tokens
+        asyncio.run(bucket.get_tokens(30))
+        
+        # Check throughput
+        throughput = bucket.get_throughput(1)
+        print(f"Throughput: {throughput:.1f} tokens/minute")
+        
+        # Enable turbo mode to bypass rate limits
+        bucket.turbo_mode_on()
+        
+        # Do some operations without rate limiting
+        asyncio.run(bucket.get_tokens(1000))  # Would normally exceed limits
+        
+        # Disable turbo mode
+        bucket.turbo_mode_off()
+        
+        # Visualize bucket usage
+        plot = bucket.visualize()
+        ```
+        """
+        pass
+    
+    # Run doctests
+    #doctest.testmod()

edsl/caching/__init__.py

@@ -0,0 +1,20 @@
+"""
+Caching module for EDSL framework.
+
+This module provides caching functionality for language model responses,
+with support for both in-memory and persistent storage through SQLite.
+It includes components for managing cache entries, handling cache initialization 
+and migration, and synchronizing with remote caches.
+
+Key components:
+- Cache: Central class for storing and retrieving language model responses
+- CacheEntry: Represents individual cached responses with metadata
+- CacheHandler: Manages cache initialization and migration
+- SQLiteDict: Dictionary-like interface to SQLite database
+"""
+
+from .cache import Cache
+from .cache_entry import CacheEntry
+from .cache_handler import CacheHandler
+
+__all__ = ["Cache", "CacheEntry", "CacheHandler"]