40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
|
"""Minimal in-memory rate limiter (per (ip, device_id), rolling 1h window).
|
||
|
|
|
||
|
|
Mirrors hesu's rate_limit.py shape. Only the chat message route uses this;
|
||
|
|
session/history endpoints are cheap and not throttled. State is process-local
|
||
|
|
and resets on restart, which is acceptable for a single-worker ai-kefu deploy.
|
||
|
|
|
||
|
|
Caveats:
|
||
|
|
- Single-process, in-memory. Under `uvicorn --workers >1` each worker keeps
|
||
|
|
its own `_buckets`, so the effective cap is `max_requests * workers`.
|
||
|
|
- Keyed on (ip, device_id); a client rotating the `X-Device-Id` header per
|
||
|
|
request bypasses it. This is a light safeguard against accidental GLM-cost
|
||
|
|
burn, NOT abuse-proof.
|
||
|
|
"""
|
||
|
|
import time
|
||
|
|
from collections import defaultdict
|
||
|
|
|
||
|
|
_buckets: dict[tuple[str, str], list[float]] = defaultdict(list)
|
||
|
|
WINDOW = 3600
|
||
|
|
DEFAULT_MAX = 30
|
||
|
|
|
||
|
|
|
||
|
|
def check(
|
||
|
|
ip: str | None,
|
||
|
|
device_id: str | None,
|
||
|
|
*,
|
||
|
|
max_requests: int = DEFAULT_MAX,
|
||
|
|
) -> bool:
|
||
|
|
"""Return True if the request is under the cap (and record it), else False."""
|
||
|
|
key = (ip or "", device_id or "")
|
||
|
|
now = time.time()
|
||
|
|
_buckets[key] = [t for t in _buckets[key] if t > now - WINDOW]
|
||
|
|
if len(_buckets[key]) >= max_requests:
|
||
|
|
return False
|
||
|
|
_buckets[key].append(now)
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def _reset_for_test() -> None:
|
||
|
|
_buckets.clear()
|