Spaces:
Running
Running
| """Author RAG — Rate Limiter Service. | |
| Redis sliding window rate limiter with in-memory fallback. | |
| Used to protect chat and auth endpoints. | |
| Per implementation plan — Fail-safe: Redis primary, in-memory fallback. | |
| """ | |
| import time | |
| from collections import defaultdict | |
| import structlog | |
| logger = structlog.get_logger(__name__) | |
| # In-memory fallback storage: {key: [(timestamp, count)]} | |
| _memory_store: dict[str, list[float]] = defaultdict(list) | |
| async def check_rate_limit( | |
| redis, | |
| key: str, | |
| max_requests: int, | |
| window_seconds: int = 60, | |
| ) -> tuple[bool, int]: | |
| """Check if a key has exceeded its rate limit using a sliding window. | |
| Tries Redis first; falls back to in-memory if Redis is unavailable. | |
| Args: | |
| redis: Redis async client (or None). | |
| key: Unique rate limit key (e.g. "chat:ip:1.2.3.4"). | |
| max_requests: Maximum requests allowed in the window. | |
| window_seconds: Window size in seconds. | |
| Returns: | |
| Tuple of (allowed: bool, remaining: int). | |
| """ | |
| try: | |
| if redis is not None: | |
| return await _redis_sliding_window(redis, key, max_requests, window_seconds) | |
| except Exception as e: | |
| logger.warning("Redis rate limit failed, using in-memory fallback", error=str(e)) | |
| return _memory_sliding_window(key, max_requests, window_seconds) | |
| async def _redis_sliding_window( | |
| redis, | |
| key: str, | |
| max_requests: int, | |
| window_seconds: int, | |
| ) -> tuple[bool, int]: | |
| """Redis ZSET sliding window rate limiter.""" | |
| now = time.time() | |
| window_start = now - window_seconds | |
| pipe = redis.pipeline() | |
| pipe.zremrangebyscore(key, 0, window_start) | |
| pipe.zadd(key, {str(now): now}) | |
| pipe.zcard(key) | |
| pipe.expire(key, window_seconds + 1) | |
| results = await pipe.execute() | |
| count = results[2] | |
| allowed = count <= max_requests | |
| remaining = max(0, max_requests - count) | |
| return allowed, remaining | |
| def _memory_sliding_window( | |
| key: str, | |
| max_requests: int, | |
| window_seconds: int, | |
| ) -> tuple[bool, int]: | |
| """In-memory sliding window fallback.""" | |
| now = time.time() | |
| window_start = now - window_seconds | |
| timestamps = _memory_store[key] | |
| # Remove expired timestamps | |
| _memory_store[key] = [ts for ts in timestamps if ts > window_start] | |
| _memory_store[key].append(now) | |
| count = len(_memory_store[key]) | |
| allowed = count <= max_requests | |
| remaining = max(0, max_requests - count) | |
| return allowed, remaining | |