import asyncio import hashlib import os import re import time from typing import Any, Dict, List, Optional from fastapi import HTTPException, Request from helper.subscriptions import ( TIER_CONFIG, USAGE_PERIODS, fetch_subscription, normalize_plan_key, usage_locks, usage_store, ) # ------------------------------------------------------------------- # Configuration / Constants # ------------------------------------------------------------------- IDENTITY_CACHE_TTL_SECONDS = 60 CLIENT_BIND_TTL_SECONDS = int( os.getenv("CLIENT_BIND_TTL_SECONDS", str(8 * 24 * 60 * 60)) ) MAX_CLIENT_ID_LENGTH = 128 MAX_CHAT_PROMPT_CHARS = int(os.getenv("MAX_CHAT_PROMPT_CHARS", "120000")) MAX_CHAT_PROMPT_BYTES = int(os.getenv("MAX_CHAT_PROMPT_BYTES", "500000")) MAX_GROQ_PROMPT_CHARS = int(os.getenv("MAX_GROQ_PROMPT_CHARS", "90000")) MAX_GROQ_PROMPT_BYTES = int(os.getenv("MAX_GROQ_PROMPT_BYTES", "350000")) MAX_MEDIA_PROMPT_CHARS = int(os.getenv("MAX_MEDIA_PROMPT_CHARS", "4000")) MAX_MEDIA_PROMPT_BYTES = int(os.getenv("MAX_MEDIA_PROMPT_BYTES", "16000")) # ------------------------------------------------------------------- # In-memory state # ------------------------------------------------------------------- identity_cache: Dict[str, Dict[str, Any]] = {} client_subject_bindings: Dict[str, Dict[str, Any]] = {} # ------------------------------------------------------------------- # Public Rate Limit APIs # ------------------------------------------------------------------- async def resolve_rate_limit_identity( request: Request, authorization: Optional[str], client_id: Optional[str] = None, ) -> tuple[str, str]: now = time.time() normalized_client_id = sanitize_client_id(client_id) default_subject = build_default_subject(request, normalized_client_id) if not authorization or not authorization.startswith("Bearer "): return "free", resolve_bound_subject(normalized_client_id, default_subject) token = authorization.split(" ", 1)[1].strip() if not token: return "free", resolve_bound_subject(normalized_client_id, default_subject) cached = identity_cache.get(token) if cached and cached.get("expires_at", 0) > now: plan_key = cached.get("plan_key", "free") subject = cached.get("subject", default_subject) bind_client_subject(normalized_client_id, subject, plan_key) return plan_key, subject try: sub = await fetch_subscription(token) except Exception: return "free", resolve_bound_subject(normalized_client_id, default_subject) if not isinstance(sub, dict) or sub.get("error"): return "free", resolve_bound_subject(normalized_client_id, default_subject) email = sub.get("email") if isinstance(email, str) and email.strip(): subject = f"user:{email.strip().lower()}" else: subject = default_subject plan_key = normalize_plan_key(sub.get("plan_key")) identity_cache[token] = { "plan_key": plan_key, "subject": subject, "expires_at": now + IDENTITY_CACHE_TTL_SECONDS, } bind_client_subject(normalized_client_id, subject, plan_key) return plan_key, subject async def enforce_rate_limit( request: Request, authorization: Optional[str], metric: str, client_id: Optional[str] = None, ) -> Dict[str, Optional[int | str]]: if metric not in usage_store: raise HTTPException(status_code=500, detail=f"Unknown limit metric: {metric}") plan_key, subject = await resolve_rate_limit_identity( request, authorization, client_id ) plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"] plan_limits = plan.get("limits", {}) limit = plan_limits.get(metric) window_key = get_usage_period_key(metric) lock = get_usage_lock(metric, subject) async with lock: bucket = usage_store[metric] entry = bucket.get(subject) if not entry or entry.get("window") != window_key: entry = {"window": window_key, "count": 0} bucket[subject] = entry if limit is not None and entry["count"] >= int(limit): raise HTTPException( status_code=429, detail=f"{metric} limit reached for {plan.get('name', 'current plan')}", ) entry["count"] += 1 remaining = None if limit is None else max(0, int(limit) - entry["count"]) return { "plan_key": plan_key, "remaining": remaining, "used": entry["count"], "window": window_key, } async def check_audio_rate_limit( request: Request, authorization: Optional[str], client_id: Optional[str] = None, ): await enforce_rate_limit(request, authorization, "audioWeekly", client_id) async def check_image_rate_limit( request: Request, authorization: Optional[str], client_id: Optional[str] = None, ): await enforce_rate_limit(request, authorization, "imagesDaily", client_id) async def check_video_rate_limit( request: Request, authorization: Optional[str], client_id: Optional[str] = None, ): await enforce_rate_limit(request, authorization, "videosDaily", client_id) # ------------------------------------------------------------------- # Prompt Utilities # ------------------------------------------------------------------- def normalize_prompt_value(prompt: Optional[str], field_name: str = "prompt") -> str: if not isinstance(prompt, str): raise HTTPException(status_code=400, detail=f"{field_name} is required") normalized = prompt.strip() if not normalized: raise HTTPException(status_code=400, detail=f"{field_name} is required") return normalized def enforce_prompt_size(prompt: str, max_chars: int, max_bytes: int, context: str): char_len = len(prompt) byte_len = len(prompt.encode("utf-8")) if char_len > max_chars or byte_len > max_bytes: raise HTTPException( status_code=413, detail=( f"{context} is too large ({char_len} chars, {byte_len} bytes). " f"Max allowed is {max_chars} chars or {max_bytes} bytes." ), ) def calculate_messages_size(messages: list) -> tuple[int, int]: total_chars = 0 total_bytes = 0 for message in messages: if not isinstance(message, dict): continue text = message_content_to_text(message.get("content")) if not text: continue total_chars += len(text) total_bytes += len(text.encode("utf-8")) return total_chars, total_bytes def extract_user_text(messages: list) -> str: return " ".join( message_content_to_text(m.get("content")) for m in messages if isinstance(m, dict) and m.get("role") == "user" ).lower() # ------------------------------------------------------------------- # Usage / Rate Limit Internals # ------------------------------------------------------------------- def get_usage_period_key(metric: str) -> str: now = time.gmtime() period = USAGE_PERIODS.get(metric, "daily") if period == "weekly": iso_year, iso_week, _ = time.strftime("%G %V %u", now).split(" ") return f"{iso_year}-W{iso_week}" return time.strftime("%Y-%m-%d", now) def get_usage_lock(metric: str, subject: str) -> asyncio.Lock: metric_locks = usage_locks.get(metric) if metric_locks is None: metric_locks = {} usage_locks[metric] = metric_locks lock = metric_locks.get(subject) if lock is None: lock = asyncio.Lock() metric_locks[subject] = lock return lock def get_usage_snapshot_for_subject( plan_key: str, subject: str ) -> Dict[str, Dict[str, Any]]: plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"] plan_limits = plan.get("limits", {}) snapshot: Dict[str, Dict[str, Any]] = {} for metric in usage_store.keys(): limit = plan_limits.get(metric) window_key = get_usage_period_key(metric) entry = usage_store[metric].get(subject) used = 0 if entry and entry.get("window") == window_key: used = max(0, int(entry.get("count", 0))) remaining = None if limit is None else max(0, int(limit) - used) snapshot[metric] = { "limit": limit, "used": used, "remaining": remaining, "window": window_key, "period": USAGE_PERIODS.get(metric, "daily"), } return snapshot # ------------------------------------------------------------------- # Identity / Client helpers # ------------------------------------------------------------------- def sanitize_client_id(raw_client_id: Optional[str]) -> Optional[str]: if not isinstance(raw_client_id, str): return None trimmed = raw_client_id.strip() if not trimmed or len(trimmed) > MAX_CLIENT_ID_LENGTH: return None if not re.match(r"^[A-Za-z0-9._:-]+$", trimmed): return None return trimmed def build_default_subject(request: Request, client_id: Optional[str]) -> str: if client_id: client_hash = hashlib.sha256(client_id.encode("utf-8")).hexdigest()[:24] return f"client:{client_hash}" host = request.client.host if request.client else "unknown" user_agent = request.headers.get("user-agent", "") ua_hash = ( hashlib.sha256(user_agent.encode("utf-8")).hexdigest()[:12] if user_agent else "noua" ) return f"anon:{host}:{ua_hash}" def bind_client_subject(client_id: Optional[str], subject: str, plan_key: str): if not client_id: return client_subject_bindings[client_id] = { "subject": subject, "plan_key": plan_key, "expires_at": time.time() + CLIENT_BIND_TTL_SECONDS, } def resolve_bound_subject(client_id: Optional[str], fallback_subject: str) -> str: if not client_id: return fallback_subject bound = client_subject_bindings.get(client_id) if not bound: return fallback_subject if bound.get("expires_at", 0) <= time.time(): client_subject_bindings.pop(client_id, None) return fallback_subject return bound.get("subject", fallback_subject) # ------------------------------------------------------------------- # Message parsing helpers # ------------------------------------------------------------------- def message_content_to_text(content: Any) -> str: if isinstance(content, str): return content if isinstance(content, list): parts: List[str] = [] for item in content: if isinstance(item, str): parts.append(item) continue if isinstance(item, dict): text = item.get("text") if isinstance(text, str): parts.append(text) return " ".join(parts) return ""