Rhinox13/chatapi

1from __future__ import annotations
2
3from collections import deque
4from dataclasses import dataclass, field
5from threading import Lock
6from time import monotonic
7
8
9@dataclass
10class MessageRateLimiter:
11    default_limit: int = 0
12    window_seconds: float = 60.0
13    _events: dict[str, deque[float]] = field(default_factory=dict)
14    _lock: Lock = field(default_factory=Lock)
15
16    def allow(self, key: str, limit: int | None = None) -> bool:
17        effective_limit = self.default_limit if limit is None else int(limit)
18        if effective_limit <= 0:
19            return True
20
21        now = monotonic()
22        with self._lock:
23            bucket = self._events.setdefault(key, deque())
24            cutoff = now - self.window_seconds
25            while bucket and bucket[0] <= cutoff:
26                bucket.popleft()
27            if len(bucket) >= effective_limit:
28                return False
29            bucket.append(now)
30            if not bucket:
31                self._events.pop(key, None)
32            return True
33