Spaces:
Runtime error
Runtime error
File size: 3,752 Bytes
a5ddf46 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """
Rate limiting utility for the Personalization API.
Implements an in-memory sliding window rate limiter.
"""
import time
from dataclasses import dataclass, field
from typing import Dict, Optional, Tuple
from config import Config
@dataclass
class RateLimitEntry:
"""
Entry in the rate limit store tracking request count and reset time.
"""
count: int = 0
reset_time: float = 0.0
class RateLimiter:
"""
In-memory rate limiter using sliding window algorithm.
Tracks requests per user and enforces rate limits.
"""
def __init__(
self,
max_requests: Optional[int] = None,
window_seconds: Optional[int] = None
):
"""
Initialize rate limiter.
Args:
max_requests: Maximum requests per window (default from config)
window_seconds: Window duration in seconds (default from config)
"""
self.max_requests = max_requests or Config.PERSONALIZE_RATE_LIMIT
self.window_seconds = window_seconds or Config.PERSONALIZE_RATE_LIMIT_WINDOW
self._store: Dict[str, RateLimitEntry] = {}
def check(self, user_id: str) -> Tuple[bool, Optional[int]]:
"""
Check if a request is allowed for the given user.
Args:
user_id: Unique identifier for the user
Returns:
Tuple of (is_allowed, retry_after_seconds)
- If allowed: (True, None)
- If rate limited: (False, seconds_until_reset)
"""
now = time.time()
entry = self._store.get(user_id)
# No entry or window expired - create new window
if not entry or now > entry.reset_time:
self._store[user_id] = RateLimitEntry(
count=1,
reset_time=now + self.window_seconds
)
return (True, None)
# Check if over limit
if entry.count >= self.max_requests:
retry_after = int(entry.reset_time - now) + 1 # Add 1 to ensure we're past reset
return (False, retry_after)
# Increment count
entry.count += 1
return (True, None)
def get_remaining(self, user_id: str) -> int:
"""
Get remaining requests for the user in current window.
Args:
user_id: Unique identifier for the user
Returns:
Number of remaining requests
"""
now = time.time()
entry = self._store.get(user_id)
if not entry or now > entry.reset_time:
return self.max_requests
return max(0, self.max_requests - entry.count)
def reset(self, user_id: str) -> None:
"""
Reset rate limit for a user (useful for testing).
Args:
user_id: Unique identifier for the user
"""
if user_id in self._store:
del self._store[user_id]
def clear_expired(self) -> int:
"""
Clear all expired entries from the store.
Returns:
Number of entries cleared
"""
now = time.time()
expired_keys = [
key for key, entry in self._store.items()
if now > entry.reset_time
]
for key in expired_keys:
del self._store[key]
return len(expired_keys)
# Global rate limiter instance for personalization endpoint
_personalize_rate_limiter: Optional[RateLimiter] = None
def get_personalize_rate_limiter() -> RateLimiter:
"""
Get or create the global rate limiter instance.
Returns:
RateLimiter instance
"""
global _personalize_rate_limiter
if _personalize_rate_limiter is None:
_personalize_rate_limiter = RateLimiter()
return _personalize_rate_limiter
|