File size: 1,928 Bytes
1b985c6 fc29964 1b985c6 fc29964 1b985c6 fc29964 1b985c6 fc29964 1b985c6 fc29964 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | from __future__ import annotations
import threading
from datetime import datetime, timezone
from typing import Dict, Tuple, Any
from .config import settings
_lock = threading.Lock()
_state: Dict[str, Any] = {"date": None, "count": 0}
def _today_key() -> str:
return datetime.now(timezone.utc).strftime("%Y%m%d")
def _in_memory_increment(limit: int) -> Tuple[bool, Dict[str, Any]]:
info: Dict[str, Any] = {"limit": limit}
with _lock:
today = _today_key()
if _state["date"] != today:
_state["date"] = today
_state["count"] = 0
_state["count"] += 1
count = int(_state["count"])
info["count"] = count
info["remaining"] = max(0, limit - count)
return count <= limit, info
def _redis_increment(limit: int) -> Tuple[bool, Dict[str, Any]]:
info: Dict[str, Any] = {"limit": limit}
import redis # local import to keep optional
client = redis.from_url(settings.redis_url, decode_responses=True)
key = f"toxrai:ai_cap:{_today_key()}"
count = int(client.incr(key))
if count == 1:
# expire a bit after 24h to avoid unbounded growth
client.expire(key, 60 * 60 * 48)
info["count"] = count
info["remaining"] = max(0, limit - count)
return count <= limit, info
def check_and_increment_global_ai_cap() -> Tuple[bool, Dict[str, Any]]:
"""Global daily cap to prevent overuse in public demos."""
limit = int(settings.max_ai_summaries_per_day or 0)
if limit <= 0:
return True, {"limit": limit, "count": 0, "remaining": None}
if settings.redis_url:
try:
return _redis_increment(limit)
except Exception as e:
# Fall back to in-memory, but report the error for debugging.
allowed, info = _in_memory_increment(limit)
info["redis_error"] = str(e)
return allowed, info
return _in_memory_increment(limit)
|