#!/usr/bin/env python3 """Public review service for the WorldModelBench Hugging Face Space.""" from __future__ import annotations import io import json import os import threading import uuid from contextlib import closing from datetime import datetime, timezone from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from urllib.error import HTTPError, URLError from urllib.parse import parse_qs, urlparse from urllib.request import Request, urlopen from huggingface_hub import CommitOperationAdd, HfApi, get_token ROOT = Path(__file__).resolve().parent STATIC_ROOT = ROOT / "static" MEDIA_CACHE_ROOT = ROOT / "media-cache" CATALOG_PATH = STATIC_ROOT / "data" / "worldmodelbench-catalog.js" DATASET_REPO = os.environ.get("WORLDMODELBENCH_REPO", "Xinyun-Liu/WorldModelBench") HF_TOKEN = os.environ.get("HF_TOKEN") or get_token() or "" MAX_BODY_BYTES = 32_000 MIN_SUBMISSION_INTERVAL = 0.8 last_submission_by_ip: dict[str, float] = {} rate_lock = threading.Lock() def read_catalog() -> dict[str, dict[str, str]]: raw = CATALOG_PATH.read_text(encoding="utf-8") payload = raw.split("=", 1)[1].strip().rstrip(";") rows = json.loads(payload) return {row["videoPath"]: row for row in rows if isinstance(row, dict) and row.get("videoPath")} CATALOG = read_catalog() def cached_media_path(video_path: str) -> Path | None: """Return a cache entry only when it stays beneath the cache root.""" cache_root = MEDIA_CACHE_ROOT.resolve() try: candidate = (cache_root / video_path).resolve() candidate.relative_to(cache_root) except ValueError: return None return candidate if candidate.is_file() else None def utc_now() -> str: return datetime.now(timezone.utc).isoformat(timespec="seconds") def clean_string(value: object, field: str, maximum: int, required: bool = False) -> str: if not isinstance(value, str): raise ValueError(f"{field}_must_be_text") value = value.strip() if required and not value: raise ValueError(f"{field}_is_required") if len(value) > maximum: raise ValueError(f"{field}_is_too_long") return value def validate_submission(payload: object) -> dict[str, str]: if not isinstance(payload, dict): raise ValueError("request_must_be_a_json_object") annotator = clean_string(payload.get("annotator"), "annotator", 80, required=True) video_path = clean_string(payload.get("video_path"), "video_path", 500, required=True) row = CATALOG.get(video_path) if row is None: raise ValueError("video_path_is_not_in_the_catalog") decision = clean_string(payload.get("decision", "pending"), "decision", 20, required=True) if decision not in {"pending", "keep", "reject"}: raise ValueError("invalid_decision") return { "annotator": annotator, "video_path": video_path, "sample_id": clean_string(row.get("id", ""), "sample_id", 300, required=True), "data_source": clean_string(row.get("sourceDataset", ""), "data_source", 200, required=True), "decision": decision, "track": clean_string(payload.get("track", ""), "track", 200), "scene": clean_string(payload.get("scene", ""), "scene", 300), "spatial_ability": clean_string(payload.get("spatial_ability", ""), "spatial_ability", 500), "perspective": clean_string(payload.get("perspective", ""), "perspective", 300), "text_caption": clean_string(payload.get("text_caption", ""), "text_caption", 4000), } class ReviewHandler(SimpleHTTPRequestHandler): server_version = "WorldModelBenchReview/1.0" def __init__(self, *args, **kwargs): super().__init__(*args, directory=str(STATIC_ROOT), **kwargs) def end_headers(self) -> None: if urlparse(self.path).path == "/review.html": self.send_header("Cache-Control", "no-store") super().end_headers() def send_json(self, status: int, value: dict[str, object]) -> None: body = json.dumps(value, ensure_ascii=False).encode("utf-8") self.send_response(status) self.send_header("Content-Type", "application/json; charset=utf-8") self.send_header("Content-Length", str(len(body))) self.send_header("Cache-Control", "no-store") self.end_headers() self.wfile.write(body) def proxy_asset(self, url: str, send_body: bool = True) -> None: """Relay an HF resolve URL as a same-origin, range-capable response.""" headers = {"User-Agent": "WorldModelBench-Review/1.0"} requested_range = self.headers.get("Range") if requested_range: headers["Range"] = requested_range try: with closing(urlopen(Request(url, headers=headers), timeout=90)) as upstream: self.send_response(upstream.status) for name in ("Content-Type", "Content-Length", "Content-Range", "Accept-Ranges", "Last-Modified", "ETag"): value = upstream.headers.get(name) if value: self.send_header(name, value) self.send_header("Cache-Control", "public, max-age=300") self.end_headers() if send_body: while chunk := upstream.read(1024 * 1024): self.wfile.write(chunk) except HTTPError as error: self.send_json(error.code, {"error": "upstream_asset_unavailable"}) except (OSError, URLError, TimeoutError): self.send_json(502, {"error": "upstream_asset_connection_failed"}) def serve_cached_media(self, asset: Path, send_body: bool = True) -> None: """Serve a locally cached original MP4 with browser Range support.""" total_size = asset.stat().st_size start, end, partial = 0, total_size - 1, False requested_range = self.headers.get("Range", "") if requested_range: try: unit, spec = requested_range.split("=", 1) if unit.strip() != "bytes" or "," in spec: raise ValueError first, last = spec.strip().split("-", 1) if first: start = int(first) end = int(last) if last else total_size - 1 elif last: suffix = int(last) if suffix <= 0: raise ValueError start = max(total_size - suffix, 0) else: raise ValueError if start < 0 or start >= total_size or end < start: raise ValueError end = min(end, total_size - 1) partial = True except (ValueError, TypeError): self.send_response(416) self.send_header("Content-Range", f"bytes */{total_size}") self.send_header("Content-Length", "0") self.end_headers() return length = end - start + 1 self.send_response(206 if partial else 200) self.send_header("Content-Type", "video/mp4") self.send_header("Content-Length", str(length)) self.send_header("Accept-Ranges", "bytes") if partial: self.send_header("Content-Range", f"bytes {start}-{end}/{total_size}") self.send_header("Cache-Control", "public, max-age=300") self.end_headers() if not send_body: return try: with asset.open("rb") as stream: stream.seek(start) remaining = length while remaining: chunk = stream.read(min(1024 * 1024, remaining)) if not chunk: break self.wfile.write(chunk) remaining -= len(chunk) except BrokenPipeError: return def proxy_catalog_asset(self, kind: str, query: dict[str, list[str]], send_body: bool = True) -> None: paths = query.get("path", []) if len(paths) != 1: self.send_json(400, {"error": "exactly_one_video_path_is_required"}) return # Reload the small catalog at request time so a catalog update never # leaves the media proxy serving a stale playback URL. row = read_catalog().get(paths[0]) if row is None: self.send_json(404, {"error": "video_path_is_not_in_the_catalog"}) return if kind == "media": cached_asset = cached_media_path(paths[0]) # A cached source original must not override a catalog entry that # deliberately points at a browser-compatible derivative. cache_matches_catalog = isinstance(row.get("videoUrl"), str) and row["videoUrl"].endswith(f"/{paths[0]}") if cached_asset is not None and cache_matches_catalog: self.serve_cached_media(cached_asset, send_body=send_body) return url = row.get("videoUrl") if kind == "media" else row.get("initialFrameUrl") if not isinstance(url, str) or not url: self.send_json(404, {"error": "preview_is_not_available"}) return self.proxy_asset(url, send_body=send_body) def do_GET(self) -> None: parsed = urlparse(self.path) path = parsed.path if path == "/api/health": self.send_json(200 if HF_TOKEN else 503, {"ok": bool(HF_TOKEN), "catalog_videos": len(CATALOG), "dataset_repo": DATASET_REPO}) return if path == "/media": self.proxy_catalog_asset("media", parse_qs(parsed.query, keep_blank_values=True)) return if path == "/preview": self.proxy_catalog_asset("preview", parse_qs(parsed.query, keep_blank_values=True)) return if path == "/": self.path = "/review.html" super().do_GET() def do_HEAD(self) -> None: parsed = urlparse(self.path) if parsed.path == "/media": self.proxy_catalog_asset("media", parse_qs(parsed.query, keep_blank_values=True), send_body=False) return if parsed.path == "/preview": self.proxy_catalog_asset("preview", parse_qs(parsed.query, keep_blank_values=True), send_body=False) return if parsed.path == "/": self.path = "/review.html" super().do_HEAD() def do_POST(self) -> None: if urlparse(self.path).path != "/api/submissions": self.send_json(404, {"error": "not_found"}) return if not HF_TOKEN: self.send_json(503, {"error": "submission_service_not_configured"}) return client_ip = self.client_address[0] now = datetime.now(timezone.utc).timestamp() with rate_lock: previous = last_submission_by_ip.get(client_ip, 0.0) if now - previous < MIN_SUBMISSION_INTERVAL: self.send_json(429, {"error": "please_wait_before_submitting_again"}) return last_submission_by_ip[client_ip] = now try: content_length = int(self.headers.get("Content-Length", "0")) if content_length <= 0 or content_length > MAX_BODY_BYTES: raise ValueError("invalid_request_size") payload = json.loads(self.rfile.read(content_length).decode("utf-8")) record = validate_submission(payload) submission_id = uuid.uuid4().hex created_at = utc_now() record.update({ "schema_version": "worldmodelbench-web-review-v1", "submission_id": submission_id, "submitted_at": created_at, "source": "public_review_space", }) path_in_repo = f"web-review/submissions/{created_at[:10]}/{submission_id}.json" content = json.dumps(record, ensure_ascii=False, indent=2).encode("utf-8") HfApi(token=HF_TOKEN).create_commit( repo_id=DATASET_REPO, repo_type="dataset", operations=[CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(content))], commit_message=f"Add web review annotation {submission_id[:12]}", ) self.send_json(201, {"ok": True, "submission_id": submission_id, "submitted_at": created_at}) except ValueError as error: self.send_json(400, {"error": str(error)}) except Exception: self.send_json(502, {"error": "huggingface_submission_failed"}) def main() -> None: port = int(os.environ.get("PORT", "7860")) server = ThreadingHTTPServer(("0.0.0.0", port), ReviewHandler) print(f"WorldModelBench review service listening on {port}", flush=True) server.serve_forever() if __name__ == "__main__": main()