""" DAARION Platform - Rate Limiter Protects agents from spam and abuse """ import time from collections import defaultdict from typing import Dict, Tuple, Optional import structlog logger = structlog.get_logger() class RateLimiter: """Token bucket rate limiter per user/agent""" def __init__(self, default_rpm: int = 30, default_burst: int = 5): self.default_rpm = default_rpm self.default_burst = default_burst # {user_id: {agent_id: (tokens, last_update)}} self._buckets: Dict[str, Dict[str, Tuple[float, float]]] = defaultdict(dict) # {agent_id: rpm_limit} self._agent_limits: Dict[str, int] = {} # Blocked users (temporary) self._blocked: Dict[str, float] = {} # user_id: unblock_time def set_agent_limit(self, agent_id: str, rpm: int): """Set rate limit for specific agent""" self._agent_limits[agent_id] = rpm logger.info("rate_limit_set", agent_id=agent_id, rpm=rpm) def _get_limit(self, agent_id: str) -> int: """Get RPM limit for agent""" return self._agent_limits.get(agent_id, self.default_rpm) def _refill_tokens(self, user_id: str, agent_id: str) -> float: """Refill tokens based on time elapsed""" now = time.time() if agent_id not in self._buckets[user_id]: # New user - full bucket self._buckets[user_id][agent_id] = (self.default_burst, now) return self.default_burst tokens, last_update = self._buckets[user_id][agent_id] rpm = self._get_limit(agent_id) # Calculate tokens to add (tokens per second = rpm / 60) elapsed = now - last_update tokens_to_add = elapsed * (rpm / 60.0) # Cap at burst limit new_tokens = min(self.default_burst, tokens + tokens_to_add) self._buckets[user_id][agent_id] = (new_tokens, now) return new_tokens def is_blocked(self, user_id: str) -> bool: """Check if user is temporarily blocked""" if user_id not in self._blocked: return False if time.time() > self._blocked[user_id]: del self._blocked[user_id] return False return True def block_user(self, user_id: str, seconds: int = 60): """Temporarily block user""" self._blocked[user_id] = time.time() + seconds logger.warning("user_blocked", user_id=user_id, seconds=seconds) def check(self, user_id: str, agent_id: str) -> Tuple[bool, Optional[str]]: """ Check if request is allowed. Returns: (allowed, error_message) """ # Check block list if self.is_blocked(user_id): remaining = int(self._blocked[user_id] - time.time()) return False, f"Занадто багато запитів. Спробуйте через {remaining} сек." # Refill and check tokens tokens = self._refill_tokens(user_id, agent_id) if tokens < 1: # No tokens - calculate wait time rpm = self._get_limit(agent_id) wait_seconds = int(60 / rpm) logger.warning("rate_limit_exceeded", user_id=user_id, agent_id=agent_id) return False, f"Ліміт запитів. Спробуйте через {wait_seconds} сек." # Consume token self._buckets[user_id][agent_id] = (tokens - 1, time.time()) return True, None def get_remaining(self, user_id: str, agent_id: str) -> int: """Get remaining requests for user""" tokens = self._refill_tokens(user_id, agent_id) return int(tokens) # Singleton instance _limiter = None def get_rate_limiter() -> RateLimiter: """Get or create rate limiter instance""" global _limiter if _limiter is None: _limiter = RateLimiter() # Set agent-specific limits _limiter.set_agent_limit("helion", 60) _limiter.set_agent_limit("nutra", 30) _limiter.set_agent_limit("greenfood", 20) _limiter.set_agent_limit("druid", 20) _limiter.set_agent_limit("daarwizz", 100) return _limiter