| """ |
| Per-token usage tracking for assignment tokens. |
| |
| Tracks how many runs each assignment token has used and enforces |
| a maximum. A "run" = one user-triggered generation job (cognitive |
| interview, expert review, QAS, or silicon sampling). |
| |
| Storage: JSON file, either local or backed by a HuggingFace dataset |
| repo for persistence across HF Space restarts. |
| |
| Usage in dashboard.py — see integration instructions at bottom of file. |
| """ |
|
|
| import csv |
| import hashlib |
| import json |
| import logging |
| import os |
| import threading |
| from pathlib import Path |
|
|
| logger = logging.getLogger(__name__) |
|
|
| |
| |
| |
|
|
| MAX_RUNS_PER_TOKEN = 8 |
|
|
| |
| _TOKEN_OVERRIDES = { |
| "DE9BBFB5": 15, |
| } |
|
|
|
|
| def get_max_runs(token: str) -> int: |
| """Return the max runs for a token, checking overrides first.""" |
| return _TOKEN_OVERRIDES.get(token.strip(), MAX_RUNS_PER_TOKEN) |
|
|
| |
| _TOKEN_CSV_PATH = Path(__file__).parent / "student_tokens.csv" |
|
|
| |
| _LOCAL_USAGE_PATH = Path(__file__).parent / "token_usage.json" |
|
|
| |
| |
| |
| |
| _HF_USAGE_REPO = os.environ.get("HF_USAGE_REPO", "") |
| _HF_USAGE_TOKEN = os.environ.get("HF_USAGE_TOKEN", "") |
| _HF_USAGE_FILENAME = "token_usage.json" |
|
|
| |
| _lock = threading.Lock() |
|
|
|
|
| |
| |
| |
|
|
| def mask_token(token: str) -> str: |
| """Return a short hash for logging. Never log raw tokens.""" |
| return hashlib.sha256(token.encode()).hexdigest()[:8] |
|
|
|
|
| |
| |
| |
|
|
| _valid_tokens: set[str] | None = None |
|
|
|
|
| def _load_valid_tokens() -> set[str]: |
| """Load the set of valid assignment tokens from the CSV.""" |
| global _valid_tokens |
| if _valid_tokens is not None: |
| return _valid_tokens |
|
|
| tokens = set() |
| path = _TOKEN_CSV_PATH |
| if path.exists(): |
| with open(path, newline="") as f: |
| reader = csv.DictReader(f) |
| for row in reader: |
| t = row.get("token", "").strip() |
| if t: |
| tokens.add(t) |
| logger.info(f"Loaded {len(tokens)} valid assignment tokens") |
| else: |
| logger.warning(f"Token CSV not found: {path}") |
|
|
| _valid_tokens = tokens |
| return _valid_tokens |
|
|
|
|
| def is_valid_token(token: str) -> bool: |
| """Check whether a token is in the valid set.""" |
| return token.strip() in _load_valid_tokens() |
|
|
|
|
| |
| |
| |
|
|
| def _load_local() -> dict: |
| if _LOCAL_USAGE_PATH.exists(): |
| try: |
| with open(_LOCAL_USAGE_PATH) as f: |
| return json.load(f) |
| except (json.JSONDecodeError, OSError): |
| logger.warning("Corrupt local usage file — starting fresh") |
| return {} |
|
|
|
|
| def _save_local(data: dict): |
| with open(_LOCAL_USAGE_PATH, "w") as f: |
| json.dump(data, f, indent=2) |
|
|
|
|
| |
| |
| |
|
|
| def _hf_enabled() -> bool: |
| return bool(_HF_USAGE_REPO and _HF_USAGE_TOKEN) |
|
|
|
|
| def _load_hf() -> dict: |
| """Download token_usage.json from the HF dataset repo.""" |
| try: |
| from huggingface_hub import hf_hub_download |
| path = hf_hub_download( |
| repo_id=_HF_USAGE_REPO, |
| filename=_HF_USAGE_FILENAME, |
| repo_type="dataset", |
| token=_HF_USAGE_TOKEN, |
| ) |
| with open(path) as f: |
| return json.load(f) |
| except Exception as e: |
| logger.warning(f"Could not load usage from HF repo: {e}") |
| return {} |
|
|
|
|
| def _save_hf(data: dict): |
| """Upload token_usage.json to the HF dataset repo.""" |
| try: |
| from huggingface_hub import HfApi |
| import tempfile |
|
|
| api = HfApi(token=_HF_USAGE_TOKEN) |
| with tempfile.NamedTemporaryFile( |
| mode="w", suffix=".json", delete=False |
| ) as tmp: |
| json.dump(data, tmp, indent=2) |
| tmp_path = tmp.name |
|
|
| api.upload_file( |
| path_or_fileobj=tmp_path, |
| path_in_repo=_HF_USAGE_FILENAME, |
| repo_id=_HF_USAGE_REPO, |
| repo_type="dataset", |
| commit_message="Update token usage", |
| ) |
| os.unlink(tmp_path) |
| except Exception as e: |
| logger.error(f"Could not save usage to HF repo: {e}") |
|
|
|
|
| |
| |
| |
|
|
| def load_token_usage() -> dict: |
| """Load usage data. Returns dict like {"TOKEN": {"runs_used": 3}, ...}.""" |
| if _hf_enabled(): |
| return _load_hf() |
| return _load_local() |
|
|
|
|
| def save_token_usage(data: dict): |
| """Persist usage data.""" |
| _save_local(data) |
| if _hf_enabled(): |
| _save_hf(data) |
|
|
|
|
| |
| |
| |
|
|
| def get_runs_used(token: str) -> int: |
| """Return how many runs this token has consumed.""" |
| with _lock: |
| data = load_token_usage() |
| return data.get(token, {}).get("runs_used", 0) |
|
|
|
|
| def get_remaining_runs(token: str) -> int: |
| """Return how many runs this token has left.""" |
| return get_max_runs(token) - get_runs_used(token) |
|
|
|
|
| def can_use_token(token: str) -> tuple[bool, str]: |
| """ |
| Validate a token and check its quota. |
| |
| Returns (ok, message): |
| (True, "N runs remaining") — token is valid and has quota |
| (False, reason) — token is invalid or exhausted |
| """ |
| token = token.strip() |
| if not is_valid_token(token): |
| return False, "Invalid assignment token." |
|
|
| remaining = get_remaining_runs(token) |
| if remaining <= 0: |
| return False, ( |
| "You have reached the maximum number of runs allowed for this " |
| "assignment token. If you encountered a technical problem, " |
| "contact Patrick Sturgis." |
| ) |
| return True, f"{remaining} runs remaining" |
|
|
|
|
| def increment_token_usage(token: str): |
| """Increment the run count for a token. Persists immediately.""" |
| with _lock: |
| data = load_token_usage() |
| entry = data.setdefault(token, {"runs_used": 0}) |
| entry["runs_used"] += 1 |
| save_token_usage(data) |
| logger.info( |
| f"Token {mask_token(token)}: runs_used -> {entry['runs_used']}" |
| ) |
|
|
|
|
| def rollback_token_usage(token: str): |
| """Decrement the run count (on early failure). Persists immediately.""" |
| with _lock: |
| data = load_token_usage() |
| entry = data.get(token) |
| if entry and entry.get("runs_used", 0) > 0: |
| entry["runs_used"] -= 1 |
| save_token_usage(data) |
| logger.info( |
| f"Token {mask_token(token)}: rollback -> {entry['runs_used']}" |
| ) |
|
|
|
|
| def reset_token(token: str): |
| """Reset a token's usage to zero. For admin use.""" |
| with _lock: |
| data = load_token_usage() |
| data[token] = {"runs_used": 0} |
| save_token_usage(data) |
|
|
|
|
| def reset_all_tokens(): |
| """Reset all token usage to zero. For admin use.""" |
| with _lock: |
| save_token_usage({}) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
