| 1 | import asyncio |
| 2 | import time |
| 3 | from typing import Optional |
| 4 | |
| 5 | |
| 6 | class RateLimiter: |
| 7 | """ |
| 8 | Rate limiter to control API request frequency. |
| 9 | |
| 10 | Ensures that no more than max_requests_per_minute requests are made per minute |
| 11 | and no more than max_requests_per_day requests are made per day. |
| 12 | """ |
| 13 | |
| 14 | def __init__( |
| 15 | self, |
| 16 | max_requests_per_minute: Optional[int] = None, |
| 17 | max_requests_per_day: Optional[int] = None |
| 18 | ): |
| 19 | """ |
| 20 | Initialize the rate limiter. |
| 21 | |
| 22 | Args: |
| 23 | max_requests_per_minute: Maximum number of requests allowed per minute. |
| 24 | If None, no per-minute limit is enforced. |
| 25 | max_requests_per_day: Maximum number of requests allowed per day. |
| 26 | If None, no per-day limit is enforced. |
| 27 | """ |
| 28 | self.max_requests_per_minute = max_requests_per_minute |
| 29 | self.max_requests_per_day = max_requests_per_day |
| 30 | self.request_times = [] |
| 31 | self.lock = asyncio.Lock() |
| 32 | |
| 33 | # If per-minute rate limiting is enabled, calculate the minimum delay between requests |
| 34 | if max_requests_per_minute and max_requests_per_minute > 0: |
| 35 | self.min_delay = 60.0 / max_requests_per_minute |
| 36 | else: |
| 37 | self.min_delay = 0 |
| 38 | |
| 39 | async def acquire(self): |
| 40 | """ |
| 41 | Acquire permission to make a request. |
| 42 | |
| 43 | This method will block until it's safe to make a request according to the rate limits. |
| 44 | |
| 45 | The lock is only held while checking and recording, never while sleeping: |
| 46 | a caller waiting out a window (up to 24h for the daily limit) must not |
| 47 | block every other caller's check. After each sleep the limits are |
| 48 | re-checked, since another caller may have taken the freed slot. |
| 49 | """ |
| 50 | if not self.max_requests_per_minute and not self.max_requests_per_day: |
| 51 | # Rate limiting is disabled |
| 52 | return |
| 53 | |
| 54 | while True: |
| 55 | message = None |
| 56 | async with self.lock: |
| 57 | current_time = time.time() |
| 58 | |
| 59 | # Clean up old request times (keep requests from last 24 hours for daily limit) |
| 60 | if self.max_requests_per_day: |
| 61 | self.request_times = [t for t in self.request_times if current_time - t < 86400] |
| 62 | elif self.max_requests_per_minute: |
| 63 | self.request_times = [t for t in self.request_times if current_time - t < 60] |
| 64 | |
| 65 | wait_time = 0.0 |
| 66 | |
| 67 | # Check daily limit first |
| 68 | if self.max_requests_per_day and self.max_requests_per_day > 0: |
| 69 | daily_requests = [t for t in self.request_times if current_time - t < 86400] |
| 70 | if len(daily_requests) >= self.max_requests_per_day: |
| 71 | wait_time = 86400 - (current_time - daily_requests[0]) |
| 72 | hours = wait_time / 3600 |
| 73 | message = f"Daily rate limit reached ({self.max_requests_per_day} requests/day). Waiting {hours:.1f} hours..." |
| 74 | |
| 75 | # Check per-minute limit |
| 76 | if wait_time <= 0 and self.max_requests_per_minute and self.max_requests_per_minute > 0: |
| 77 | minute_requests = [t for t in self.request_times if current_time - t < 60] |
| 78 | if len(minute_requests) >= self.max_requests_per_minute: |
| 79 | wait_time = 60 - (current_time - minute_requests[0]) |
| 80 | message = f"Rate limit reached ({self.max_requests_per_minute} requests/min). Waiting {wait_time:.1f}s..." |
| 81 | elif self.request_times and self.min_delay > 0: |
| 82 | # Also ensure minimum delay between consecutive requests |
| 83 | time_since_last = current_time - self.request_times[-1] |
| 84 | if time_since_last < self.min_delay: |
| 85 | wait_time = self.min_delay - time_since_last |
| 86 | |
| 87 | if wait_time <= 0: |
| 88 | # Record this request |
| 89 | self.request_times.append(current_time) |
| 90 | return |
| 91 | |
| 92 | if message: |
| 93 | print(message) |
| 94 | await asyncio.sleep(wait_time) |
| 95 |