negoptimAi / backend /app /services /limits.py
samir12321's picture
Initial commit: Negoptim AI RAG chatbot (backend + frontend + deploy config)
af404c9
Raw
History Blame Contribute Delete
2.68 kB
"""
In-memory abuse-protection and cost-control primitives.
Single-instance by design — it matches the app's deployment model (one uvicorn
process on HF Spaces / locally). Nothing persists across restarts, which is
acceptable for a public demo; swap for Redis if the app ever scales out.
"""
import datetime
import time
from collections import defaultdict, deque
from typing import Any, Deque, Dict, Optional, Tuple
class SlidingWindowLimiter:
"""Allows at most `max_events` per `window_seconds` per key (e.g. client IP)."""
def __init__(self, max_events: int, window_seconds: float):
self.max_events = max_events
self.window = window_seconds
self._events: Dict[str, Deque[float]] = defaultdict(deque)
def allow(self, key: str) -> bool:
now = time.monotonic()
dq = self._events[key]
while dq and now - dq[0] > self.window:
dq.popleft()
if len(dq) >= self.max_events:
return False
dq.append(now)
# Opportunistic cleanup so the key map can't grow without bound
if len(self._events) > 5000:
stale = [k for k, v in self._events.items() if not v or now - v[-1] > self.window]
for k in stale:
del self._events[k]
return True
class TTLCache:
"""Tiny TTL + max-size cache (used for first-turn chat responses)."""
def __init__(self, ttl_seconds: float, max_entries: int):
self.ttl = ttl_seconds
self.max_entries = max_entries
self._data: Dict[str, Tuple[float, Any]] = {}
def get(self, key: str) -> Optional[Any]:
item = self._data.get(key)
if item is None:
return None
ts, value = item
if time.monotonic() - ts > self.ttl:
del self._data[key]
return None
return value
def set(self, key: str, value: Any) -> None:
if key not in self._data and len(self._data) >= self.max_entries:
oldest = min(self._data, key=lambda k: self._data[k][0])
del self._data[oldest]
self._data[key] = (time.monotonic(), value)
class DailyCounter:
"""Global counter that resets at UTC midnight (e.g. outgoing emails/day)."""
def __init__(self, limit: int):
self.limit = limit
self._day: Optional[str] = None
self._count = 0
def increment_if_allowed(self) -> bool:
today = datetime.datetime.now(datetime.timezone.utc).date().isoformat()
if today != self._day:
self._day = today
self._count = 0
if self._count >= self.limit:
return False
self._count += 1
return True