training/share_trace.py

"""Publish a downloaded agent trace to the Hugging Face Hub (Sharing is Caring).

Two-step, privacy-first flow:
  1. In the Space's Activity tab, click **⬇ Download trace (JSON)** — the trace
     stays on your device; the hosted Space holds NO Hub token.
  2. Run this CLI on your own machine, where you're logged in, to upload it to a
     public/private HF **dataset** repo for others to learn from.

Auth comes from your local Hugging Face login (same as training/export_gguf.sh):
    hf auth login                    # or: export HF_TOKEN=hf_xxxxxxxx

Usage:
    python training/share_trace.py trace.json
    python training/share_trace.py trace.json --repo me/imessage-cal-traces --public
    python training/share_trace.py trace.json --dry-run     # print plan, no network
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))  # make `server` importable
from server.trace import TRACE_SCHEMA  # noqa: E402  single source of truth for the schema id

DEFAULT_REPO = "n8mauer/imessage-cal-traces"


def _load_trace(path: Path) -> dict:
    try:
        trace = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as e:
        sys.exit(f"error: cannot read trace JSON at {path}: {e}")
    if not isinstance(trace, dict) or trace.get("schema") != TRACE_SCHEMA:
        sys.exit(
            f"error: {path} is not an iMessage-cal trace "
            f"(expected schema={TRACE_SCHEMA!r}). Download one from the Activity tab."
        )
    return trace


def _parse_args(argv=None) -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Upload an agent trace to a HF dataset repo.")
    p.add_argument("trace_path", help="path to a downloaded trace .json")
    p.add_argument("--repo", default=DEFAULT_REPO, help=f"HF dataset repo (default {DEFAULT_REPO})")
    vis = p.add_mutually_exclusive_group()
    vis.add_argument("--private", dest="private", action="store_true", default=True,
                     help="create the dataset repo private (default)")
    vis.add_argument("--public", dest="private", action="store_false",
                     help="create the dataset repo public")
    p.add_argument("--path-in-repo", default=None,
                   help="destination path inside the repo (default traces/<filename>)")
    p.add_argument("--dry-run", action="store_true",
                   help="print the planned upload and exit without any network call")
    return p.parse_args(argv)


def main(argv=None) -> None:
    args = _parse_args(argv)
    path = Path(args.trace_path)
    _load_trace(path)  # validate before doing anything

    path_in_repo = args.path_in_repo or f"traces/{path.name}"
    url = f"https://huggingface.co/datasets/{args.repo}"

    if args.dry_run:
        vis = "private" if args.private else "public"
        print(f"[dry-run] would upload {path} -> {args.repo}:{path_in_repo} ({vis})")
        print(f"[dry-run] dataset: {url}")
        return

    # Imported here (not at module top) so --dry-run and tests stay network-free.
    from huggingface_hub import HfApi

    api = HfApi()
    api.create_repo(repo_id=args.repo, repo_type="dataset", private=args.private, exist_ok=True)
    api.upload_file(
        path_or_fileobj=str(path),
        path_in_repo=path_in_repo,
        repo_id=args.repo,
        repo_type="dataset",
    )
    print(f"Uploaded {path.name} -> {url}/blob/main/{path_in_repo}")
    print(url)


if __name__ == "__main__":
    main()