import time import asyncio from helper.subscriptions import USAGE_PERIODS, usage_locks, usage_store, TIER_CONFIG, client_subject_bindings, CLIENT_BIND_TTL_SECONDS, MAX_CLIENT_ID_LENGTH from fastapi import Request, HTTPException from typing import Optional, Dict, Any, List import re import hashlib 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() 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 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 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 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) 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 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 "" 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 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