Spaces:
Sleeping
Sleeping
| """ | |
| Rate Limiter for OpenAI API to avoid hitting TPM (tokens per minute) limits. | |
| """ | |
| import time | |
| from typing import List, Tuple | |
| class RateLimiter: | |
| """Smart rate limiter that tracks token usage and sleeps only when necessary.""" | |
| def __init__(self, max_tokens_per_minute: int = 95000, buffer: int = 5000): | |
| """ | |
| Initialize the rate limiter. | |
| Args: | |
| max_tokens_per_minute: Maximum tokens allowed per minute (default: 95K for safety) | |
| buffer: Safety buffer to stay under limit (default: 5K) | |
| """ | |
| self.max_tokens = max_tokens_per_minute - buffer | |
| self.tokens_used: List[Tuple[float, int]] = [] # [(timestamp, tokens), ...] | |
| self.total_waits = 0 | |
| self.total_wait_time = 0.0 | |
| def wait_if_needed(self, tokens_needed: int) -> None: | |
| """ | |
| Check if we need to wait before making the next API call. | |
| Args: | |
| tokens_needed: Estimated tokens for the next API call | |
| """ | |
| now = time.time() | |
| # Remove tokens older than 60 seconds (sliding window) | |
| self.tokens_used = [ | |
| (ts, tok) for ts, tok in self.tokens_used | |
| if now - ts < 60 | |
| ] | |
| # Calculate tokens used in last 60 seconds | |
| tokens_in_window = sum(tok for _, tok in self.tokens_used) | |
| # If adding new request would exceed limit, wait | |
| if tokens_in_window + tokens_needed > self.max_tokens: | |
| # Calculate how long to wait | |
| oldest_timestamp = self.tokens_used[0][0] | |
| wait_time = 60 - (now - oldest_timestamp) + 1 # +1 for safety | |
| print(f"⏰ Rate limit approaching ({tokens_in_window + tokens_needed}/{self.max_tokens} tokens)") | |
| print(f" Waiting {wait_time:.1f}s for rate limit window to reset...") | |
| time.sleep(wait_time) | |
| self.total_waits += 1 | |
| self.total_wait_time += wait_time | |
| # Clear old tokens after waiting | |
| now = time.time() | |
| self.tokens_used = [ | |
| (ts, tok) for ts, tok in self.tokens_used | |
| if now - ts < 60 | |
| ] | |
| # Record this request | |
| self.tokens_used.append((now, tokens_needed)) | |
| def get_stats(self) -> dict: | |
| """Get statistics about rate limiting.""" | |
| return { | |
| 'total_waits': self.total_waits, | |
| 'total_wait_time': self.total_wait_time, | |
| 'current_window_tokens': sum(tok for _, tok in self.tokens_used) | |
| } | |
| def print_stats(self) -> None: | |
| """Print rate limiting statistics.""" | |
| stats = self.get_stats() | |
| print("\n" + "="*50) | |
| print("📊 Rate Limiter Statistics") | |
| print("="*50) | |
| print(f"Total waits: {stats['total_waits']}") | |
| print(f"Total wait time: {stats['total_wait_time']:.1f}s") | |
| print(f"Current window usage: {stats['current_window_tokens']} tokens") | |
| print("="*50 + "\n") | |
| def estimate_tokens(text: str, overhead: int = 800) -> int: | |
| """ | |
| Estimate tokens for a text string. | |
| Args: | |
| text: Input text | |
| overhead: Additional tokens for system prompts, formatting, etc. | |
| Returns: | |
| Estimated token count | |
| """ | |
| # Rough estimation: 1 token ≈ 4 characters | |
| content_tokens = len(str(text)) // 4 | |
| return content_tokens + overhead | |