mcal-ai 0.1.0__py3-none-any.whl

mcal/core/rate_limiter.py

@@ -0,0 +1,372 @@
+"""
+Rate limiting utilities for LLM API calls (Issue #39).
+
+Provides token bucket rate limiter to prevent excessive API costs
+and ensure compliance with provider rate limits.
+"""
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Optional, Dict, Any
+from collections import deque
+
+logger = logging.getLogger(__name__)
+
+
+def _utc_now() -> datetime:
+    """Get current UTC time (timezone-aware)."""
+    return datetime.now(timezone.utc)
+
+
+# =============================================================================
+# Cost Estimation
+# =============================================================================
+
+@dataclass
+class ModelPricing:
+    """Pricing information for a model (per 1M tokens)."""
+    input_cost_per_1m: float  # Cost per 1M input tokens
+    output_cost_per_1m: float  # Cost per 1M output tokens
+    
+    def estimate_cost(self, input_tokens: int, output_tokens: int) -> float:
+        """Estimate cost for given token counts."""
+        input_cost = (input_tokens / 1_000_000) * self.input_cost_per_1m
+        output_cost = (output_tokens / 1_000_000) * self.output_cost_per_1m
+        return input_cost + output_cost
+
+
+# Bedrock pricing (as of 2024 - update as needed)
+# https://aws.amazon.com/bedrock/pricing/
+BEDROCK_PRICING: Dict[str, ModelPricing] = {
+    # Llama 3.1 8B - Fast tier
+    "llama-3.1-8b": ModelPricing(input_cost_per_1m=0.22, output_cost_per_1m=0.22),
+    "llama-3.2-3b": ModelPricing(input_cost_per_1m=0.15, output_cost_per_1m=0.15),
+    # Llama 3.3 70B - Smart tier  
+    "llama-3.3-70b": ModelPricing(input_cost_per_1m=0.72, output_cost_per_1m=0.72),
+    "llama-3.1-70b": ModelPricing(input_cost_per_1m=0.72, output_cost_per_1m=0.72),
+    "llama-4-maverick": ModelPricing(input_cost_per_1m=0.50, output_cost_per_1m=1.50),
+}
+
+# OpenAI pricing
+OPENAI_PRICING: Dict[str, ModelPricing] = {
+    "gpt-4o": ModelPricing(input_cost_per_1m=2.50, output_cost_per_1m=10.00),
+    "gpt-4o-mini": ModelPricing(input_cost_per_1m=0.15, output_cost_per_1m=0.60),
+    "gpt-4-turbo": ModelPricing(input_cost_per_1m=10.00, output_cost_per_1m=30.00),
+}
+
+# Anthropic pricing
+ANTHROPIC_PRICING: Dict[str, ModelPricing] = {
+    "claude-sonnet-4-20250514": ModelPricing(input_cost_per_1m=3.00, output_cost_per_1m=15.00),
+    "claude-3-5-sonnet-20241022": ModelPricing(input_cost_per_1m=3.00, output_cost_per_1m=15.00),
+    "claude-3-haiku-20240307": ModelPricing(input_cost_per_1m=0.25, output_cost_per_1m=1.25),
+}
+
+
+def get_pricing(model: str, provider: str = "bedrock") -> Optional[ModelPricing]:
+    """Get pricing for a model."""
+    pricing_maps = {
+        "bedrock": BEDROCK_PRICING,
+        "openai": OPENAI_PRICING,
+        "anthropic": ANTHROPIC_PRICING,
+    }
+    return pricing_maps.get(provider, {}).get(model)
+
+
+# =============================================================================
+# Token Bucket Rate Limiter
+# =============================================================================
+
+@dataclass
+class RateLimitConfig:
+    """Configuration for rate limiting."""
+    requests_per_minute: int = 60  # RPM limit
+    tokens_per_minute: int = 100_000  # TPM limit
+    max_cost_per_hour: Optional[float] = None  # Cost limit (USD)
+    enable_cost_logging: bool = True  # Log cost estimates
+    warn_at_cost_percent: float = 0.8  # Warn when this % of hourly limit used
+
+
+@dataclass 
+class RateLimitStats:
+    """Statistics from rate limiter."""
+    total_requests: int = 0
+    total_input_tokens: int = 0
+    total_output_tokens: int = 0
+    total_estimated_cost: float = 0.0
+    requests_delayed: int = 0
+    total_delay_seconds: float = 0.0
+    start_time: datetime = field(default_factory=_utc_now)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        runtime = (_utc_now() - self.start_time).total_seconds()
+        return {
+            "total_requests": self.total_requests,
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+            "total_tokens": self.total_input_tokens + self.total_output_tokens,
+            "total_estimated_cost_usd": round(self.total_estimated_cost, 4),
+            "requests_delayed": self.requests_delayed,
+            "total_delay_seconds": round(self.total_delay_seconds, 2),
+            "runtime_seconds": round(runtime, 2),
+            "avg_requests_per_minute": round(self.total_requests / (runtime / 60), 2) if runtime > 0 else 0,
+        }
+
+
+class TokenBucketRateLimiter:
+    """
+    Token bucket rate limiter for LLM API calls.
+    
+    Implements both RPM (requests per minute) and TPM (tokens per minute)
+    limits using the token bucket algorithm with smooth refilling.
+    
+    Example:
+        limiter = TokenBucketRateLimiter(
+            config=RateLimitConfig(
+                requests_per_minute=60,
+                tokens_per_minute=100_000,
+                max_cost_per_hour=10.0
+            ),
+            model="llama-3.3-70b",
+            provider="bedrock"
+        )
+        
+        # Before each API call
+        await limiter.acquire(estimated_tokens=1000)
+        
+        # After each API call
+        limiter.record_usage(input_tokens=500, output_tokens=200)
+    """
+    
+    def __init__(
+        self,
+        config: Optional[RateLimitConfig] = None,
+        model: str = "llama-3.3-70b",
+        provider: str = "bedrock",
+    ):
+        self.config = config or RateLimitConfig()
+        self.model = model
+        self.provider = provider
+        self.pricing = get_pricing(model, provider)
+        
+        # Token buckets
+        self._request_bucket = float(self.config.requests_per_minute)
+        self._token_bucket = float(self.config.tokens_per_minute)
+        self._last_refill = time.monotonic()
+        
+        # Cost tracking (sliding window for hourly limit)
+        self._hourly_costs: deque = deque()  # (timestamp, cost) pairs
+        
+        # Statistics
+        self.stats = RateLimitStats()
+        
+        # Lock for thread safety
+        self._lock = asyncio.Lock()
+    
+    def _refill_buckets(self) -> None:
+        """Refill token buckets based on elapsed time."""
+        now = time.monotonic()
+        elapsed = now - self._last_refill
+        self._last_refill = now
+        
+        # Refill at rate of limit per minute
+        request_refill = (elapsed / 60.0) * self.config.requests_per_minute
+        token_refill = (elapsed / 60.0) * self.config.tokens_per_minute
+        
+        self._request_bucket = min(
+            self._request_bucket + request_refill,
+            float(self.config.requests_per_minute)
+        )
+        self._token_bucket = min(
+            self._token_bucket + token_refill,
+            float(self.config.tokens_per_minute)
+        )
+    
+    def _get_hourly_cost(self) -> float:
+        """Get total cost in the last hour."""
+        now = time.time()
+        hour_ago = now - 3600
+        
+        # Remove old entries
+        while self._hourly_costs and self._hourly_costs[0][0] < hour_ago:
+            self._hourly_costs.popleft()
+        
+        return sum(cost for _, cost in self._hourly_costs)
+    
+    def _check_cost_limit(self, estimated_cost: float) -> bool:
+        """Check if we're within the hourly cost limit."""
+        if self.config.max_cost_per_hour is None:
+            return True
+        
+        current_hourly = self._get_hourly_cost()
+        projected = current_hourly + estimated_cost
+        
+        # Warn if approaching limit
+        if self.config.enable_cost_logging:
+            usage_percent = current_hourly / self.config.max_cost_per_hour
+            if usage_percent >= self.config.warn_at_cost_percent:
+                logger.warning(
+                    f"Rate limiter: {usage_percent:.1%} of hourly cost limit used "
+                    f"(${current_hourly:.4f} / ${self.config.max_cost_per_hour:.2f})"
+                )
+        
+        return projected <= self.config.max_cost_per_hour
+    
+    async def acquire(self, estimated_tokens: int = 1000) -> float:
+        """
+        Acquire permission to make an API call.
+        
+        Blocks until rate limits allow the request, or raises if cost limit exceeded.
+        
+        Args:
+            estimated_tokens: Estimated total tokens for the request
+            
+        Returns:
+            Delay in seconds that was waited (0 if no delay)
+            
+        Raises:
+            RuntimeError: If hourly cost limit would be exceeded
+        """
+        async with self._lock:
+            total_delay = 0.0
+            
+            while True:
+                self._refill_buckets()
+                
+                # Check if we have capacity
+                if self._request_bucket >= 1.0 and self._token_bucket >= estimated_tokens:
+                    # Check cost limit
+                    if self.pricing:
+                        # Estimate cost (assume 50/50 input/output split for estimation)
+                        estimated_cost = self.pricing.estimate_cost(
+                            estimated_tokens // 2, 
+                            estimated_tokens // 2
+                        )
+                        if not self._check_cost_limit(estimated_cost):
+                            raise RuntimeError(
+                                f"Hourly cost limit (${self.config.max_cost_per_hour:.2f}) "
+                                f"would be exceeded. Current: ${self._get_hourly_cost():.4f}"
+                            )
+                    
+                    # Consume from buckets
+                    self._request_bucket -= 1.0
+                    self._token_bucket -= estimated_tokens
+                    
+                    if total_delay > 0:
+                        self.stats.requests_delayed += 1
+                        self.stats.total_delay_seconds += total_delay
+                        logger.debug(f"Rate limiter: delayed {total_delay:.2f}s")
+                    
+                    return total_delay
+                
+                # Calculate wait time
+                wait_for_request = 0.0
+                wait_for_tokens = 0.0
+                
+                if self._request_bucket < 1.0:
+                    # Time to refill 1 request
+                    wait_for_request = (1.0 - self._request_bucket) * 60.0 / self.config.requests_per_minute
+                
+                if self._token_bucket < estimated_tokens:
+                    # Time to refill needed tokens
+                    needed = estimated_tokens - self._token_bucket
+                    wait_for_tokens = needed * 60.0 / self.config.tokens_per_minute
+                
+                wait_time = max(wait_for_request, wait_for_tokens, 0.1)  # Min 100ms
+                wait_time = min(wait_time, 60.0)  # Max 60s
+                
+                logger.debug(
+                    f"Rate limiter: waiting {wait_time:.2f}s "
+                    f"(requests: {self._request_bucket:.1f}, tokens: {self._token_bucket:.0f})"
+                )
+                
+                # Release lock while waiting
+                self._lock.release()
+                try:
+                    await asyncio.sleep(wait_time)
+                    total_delay += wait_time
+                finally:
+                    await self._lock.acquire()
+    
+    def record_usage(self, input_tokens: int, output_tokens: int) -> None:
+        """
+        Record actual token usage after an API call.
+        
+        Args:
+            input_tokens: Actual input tokens used
+            output_tokens: Actual output tokens used
+        """
+        self.stats.total_requests += 1
+        self.stats.total_input_tokens += input_tokens
+        self.stats.total_output_tokens += output_tokens
+        
+        # Calculate and record cost
+        if self.pricing:
+            cost = self.pricing.estimate_cost(input_tokens, output_tokens)
+            self.stats.total_estimated_cost += cost
+            self._hourly_costs.append((time.time(), cost))
+            
+            if self.config.enable_cost_logging:
+                logger.info(
+                    f"LLM call: {input_tokens} in + {output_tokens} out tokens, "
+                    f"est. cost: ${cost:.6f}, total: ${self.stats.total_estimated_cost:.4f}"
+                )
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get rate limiter statistics."""
+        stats = self.stats.to_dict()
+        stats["model"] = self.model
+        stats["provider"] = self.provider
+        stats["config"] = {
+            "requests_per_minute": self.config.requests_per_minute,
+            "tokens_per_minute": self.config.tokens_per_minute,
+            "max_cost_per_hour": self.config.max_cost_per_hour,
+        }
+        stats["current_buckets"] = {
+            "requests": round(self._request_bucket, 2),
+            "tokens": round(self._token_bucket, 0),
+        }
+        return stats
+    
+    def reset_stats(self) -> None:
+        """Reset statistics."""
+        self.stats = RateLimitStats()
+        self._hourly_costs.clear()
+
+
+# =============================================================================
+# Convenience Functions
+# =============================================================================
+
+def create_rate_limiter(
+    model: str = "llama-3.3-70b",
+    provider: str = "bedrock",
+    rpm: int = 60,
+    tpm: int = 100_000,
+    max_cost_per_hour: Optional[float] = None,
+    enable_cost_logging: bool = True,
+) -> TokenBucketRateLimiter:
+    """
+    Create a rate limiter with the specified configuration.
+    
+    Args:
+        model: Model name for cost estimation
+        provider: Provider name ("bedrock", "openai", "anthropic")
+        rpm: Requests per minute limit
+        tpm: Tokens per minute limit
+        max_cost_per_hour: Maximum cost per hour in USD (None = unlimited)
+        enable_cost_logging: Log cost estimates for each call
+        
+    Returns:
+        Configured TokenBucketRateLimiter
+    """
+    config = RateLimitConfig(
+        requests_per_minute=rpm,
+        tokens_per_minute=tpm,
+        max_cost_per_hour=max_cost_per_hour,
+        enable_cost_logging=enable_cost_logging,
+    )
+    return TokenBucketRateLimiter(config=config, model=model, provider=provider)