Spaces:
Sleeping
Sleeping
File size: 2,854 Bytes
c44fab6 0bb4dfa c44fab6 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | from __future__ import annotations
import logging
import os
from collections import defaultdict
from datetime import date
from fastapi import Request
LOG = logging.getLogger(__name__)
ORG_NAME = os.getenv("HF_RATE_LIMIT_ORG", "neongeckocom")
# CCAI demo bumps the per-IP daily cap from LLMChats3's 20 to 30 to match
# the heavier multi-participant conversation pattern (the orchestrator-call
# backstop and the participant-message failsafe handle per-conversation
# cost).
DAILY_LIMIT = int(os.getenv("HF_RATE_LIMIT_DAILY", "30"))
_ip_counts: dict[str, dict] = defaultdict(lambda: {"date": "", "count": 0})
def _today() -> str:
return date.today().isoformat()
def is_org_member(request: Request) -> bool:
try:
from huggingface_hub import parse_huggingface_oauth
oauth_info = parse_huggingface_oauth(request)
if oauth_info is None:
return False
orgs = oauth_info.user_info.orgs or []
return any(o.preferred_username == ORG_NAME for o in orgs)
except Exception:
return False
def get_oauth_username(request: Request) -> str | None:
try:
from huggingface_hub import parse_huggingface_oauth
oauth_info = parse_huggingface_oauth(request)
if oauth_info is None:
return None
return oauth_info.user_info.preferred_username
except Exception:
return None
def get_client_ip(request: Request) -> str:
forwarded = request.headers.get("x-forwarded-for")
if forwarded:
return forwarded.split(",")[0].strip()
return request.client.host if request.client else "unknown"
def check_rate_limit(request: Request) -> tuple[bool, int]:
"""Returns (allowed, remaining). If org member, remaining is -1 (unlimited)."""
if is_org_member(request):
return True, -1
ip = get_client_ip(request)
today = _today()
entry = _ip_counts[ip]
if entry["date"] != today:
entry["date"] = today
entry["count"] = 0
remaining = max(0, DAILY_LIMIT - entry["count"])
if entry["count"] >= DAILY_LIMIT:
return False, 0
return True, remaining
def record_conversation(request: Request) -> None:
"""Increment the counter after a conversation is successfully started."""
if is_org_member(request):
return
ip = get_client_ip(request)
today = _today()
entry = _ip_counts[ip]
if entry["date"] != today:
entry["date"] = today
entry["count"] = 0
entry["count"] += 1
def get_remaining(request: Request) -> int:
"""Get remaining conversations for a client. -1 means unlimited."""
if is_org_member(request):
return -1
ip = get_client_ip(request)
today = _today()
entry = _ip_counts[ip]
if entry["date"] != today:
return DAILY_LIMIT
return max(0, DAILY_LIMIT - entry["count"])
|