build_local.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "huggingface_hub>=0.27",
#   "pandas>=2.0",
#   "pyarrow>=15",
#   "matplotlib>=3.8",
# ]
# ///
"""Build the public-safe derived view of the Hub's agent-usage telemetry.

The source dataset id is injected via the AGENT_USAGE_SRC env var (on Jobs:
`-e AGENT_USAGE_SRC=...`), so this public script never names it.

This is the scheduled HF Job body. It:
  1. fetches the private source (monthly + daily parquets)
  2. keeps only the public columns (relative shares; raw counts are never
     published) and only registered harness rows (the rest fold into `unknown`)
  3. writes derived parquets, renders the card PNGs, and pushes everything
     (data + charts + card + this script) to the public dataset repo
"""

import os
from datetime import date
from pathlib import Path

import pandas as pd
from huggingface_hub import snapshot_download
from huggingface_hub.constants import ENDPOINT
from huggingface_hub.utils import get_session

try:
    import agent_usage as au
except ModuleNotFoundError:
    # Scheduled Jobs run this script by URL, without the sibling module on disk:
    # fetch agent_usage.py from the same dataset repo the job pushes to.
    import sys

    from huggingface_hub import hf_hub_download

    _mod = hf_hub_download(
        os.environ["AGENT_USAGE_PUSH_REPO"], "agent_usage.py", repo_type="dataset"
    )
    sys.path.insert(0, str(Path(_mod).parent))
    import agent_usage as au

SRC = os.environ.get("AGENT_USAGE_SRC")
if not SRC:
    raise SystemExit("Set AGENT_USAGE_SRC to the source dataset id (on Jobs: -e AGENT_USAGE_SRC=...)")
# Allowlist: only these columns are ever published. Anything else in the source
# (whatever it is, now or later) is dropped by construction.
KEEP = {"month", "day", "agent", "pct_requests", "pct_users"}
# Publish the library-level rows only: the CLI source is a strict subset of
# huggingface_hub (the CLI imports the library), so keeping both would double-count.
PUBLISH_SOURCE = "huggingface_hub"
ROLLOUT_MONTH = "2026-04"            # skip pre-rollout monthly files (mostly empty)

# The `agent` value is extracted server-side from raw User-Agent strings, so it
# is arbitrary user-controlled input. Sanitize rows the same way KEEP sanitizes
# columns: only registered harness names are published; everything else folds
# into `unknown` (which is also what the card promises for unregistered tools).
# Fail loud if the registry is unreachable — never publish unfiltered tokens.
_resp = get_session().get(f"{ENDPOINT}/api/agent-harnesses", timeout=10)
_resp.raise_for_status()
PUBLISH_AGENTS = set(_resp.json()["harnesses"]) | {"unknown"}

HERE = Path(__file__).parent
OUT = HERE / "preview"
DATA = OUT / "data"


def derive(folder: str, key_col: str) -> pd.DataFrame:
    """Pull one granularity from the source, drop totals, write derived parquets."""
    src_root = Path(
        snapshot_download(SRC, repo_type="dataset", allow_patterns=f"data/{folder}/*.parquet")
    )
    out_dir = DATA / folder
    out_dir.mkdir(parents=True, exist_ok=True)

    frames = []
    for parquet in sorted((src_root / "data" / folder).glob("*.parquet")):
        # monthly: skip files before rollout; daily: keep all (chart filters by day)
        if folder == "monthly" and parquet.stem < ROLLOUT_MONTH:
            continue
        df = pd.read_parquet(parquet)
        if key_col not in df.columns:
            df[key_col] = parquet.stem
        if "source" in df.columns:
            df = df[df["source"] == PUBLISH_SOURCE]
        df = df[[c for c in df.columns if c in KEEP]]
        # Fold unregistered tokens into `unknown`, then merge the folded rows.
        df.loc[~df["agent"].isin(PUBLISH_AGENTS), "agent"] = "unknown"
        group_cols = [c for c in df.columns if not c.startswith("pct_")]
        df = df.groupby(group_cols, as_index=False).sum()
        df.to_parquet(out_dir / parquet.name, index=False)
        frames.append(df)

    if not frames:
        raise SystemExit(f"No {folder} parquets found")
    all_df = pd.concat(frames, ignore_index=True)

    # Fail loud if the published frame somehow holds a non-allowlisted column,
    # or if the source schema drifted away from the shares we expect.
    extras = set(all_df.columns) - KEEP
    if extras:
        raise SystemExit(f"Leak: non-allowlisted columns survived in {folder}: {extras}")
    if not {key_col, "agent", "pct_requests"} <= set(all_df.columns):
        raise SystemExit(f"Source schema drift: expected share columns missing in {folder}")

    print(f"✓ {folder}: {len(frames)} parquets → {out_dir}/  ({len(all_df)} rows)")
    return all_df


# ---------------------------- derive both granularities ----------------------

monthly = derive("monthly", "month")
daily = derive("daily", "day")

# pct sanity check (already 0-100 in source)
gs = monthly.groupby("month")["pct_requests"].sum()
if not ((gs > 95) & (gs < 105)).all():
    print(f"⚠ monthly pct_requests group sums look off: {gs.to_dict()}")

latest = au.latest_month(monthly)
leaders = au.top_agents(monthly, n=5)
print(f"\nLatest month: {latest}")
print(au.leaderboard_table(monthly).to_string(index=False))
print(f"\nTrend cohort (top 5, {latest}): {leaders}")

# ---------------------------- render PNGs (for tweet / card embedding) -------

au.save_png(au.leaderboard_figure(monthly), OUT / "leaderboard.png")
au.save_png(au.trend_figure(daily, agents=leaders), OUT / "trend.png")
print(f"\n✓ {OUT/'leaderboard.png'}")
print(f"✓ {OUT/'trend.png'}")

# ---------------------------- push to the public dataset repo ----------------
# Only when AGENT_USAGE_PUSH_REPO is set (on Jobs: -e AGENT_USAGE_PUSH_REPO=...).
# Everything above stays local otherwise. One commit per run: derived parquets,
# chart PNGs, the card, and this script + its module (the build is auditable
# from the dataset repo itself).

PUSH_REPO = os.environ.get("AGENT_USAGE_PUSH_REPO")
if PUSH_REPO:
    from huggingface_hub import CommitOperationAdd, HfApi

    api = HfApi()
    # Created private: review the rendered card/viewer, then flip to public in
    # settings. exist_ok means later scheduled runs never touch visibility.
    api.create_repo(PUSH_REPO, repo_type="dataset", private=True, exist_ok=True)

    # Idempotent: the job can run more often than the source updates (a new
    # source snapshot means a new monthly parquet). Nothing new -> no commit.
    if api.file_exists(
        PUSH_REPO, f"data/monthly/{latest}.parquet", repo_type="dataset"
    ) and not os.environ.get("AGENT_USAGE_FORCE"):
        print(
            f"\n✓ {PUSH_REPO} already has {latest} — nothing new, skipping push "
            "(AGENT_USAGE_FORCE=1 overrides)"
        )
        raise SystemExit(0)

    ops = []
    for parquet in sorted(DATA.rglob("*.parquet")):
        ops.append(CommitOperationAdd(str(parquet.relative_to(OUT)), str(parquet)))
    for png in ("leaderboard.png", "trend.png"):
        ops.append(CommitOperationAdd(f"assets/{png}", str(OUT / png)))
    # The build must stay auditable and self-contained from the dataset repo
    # itself: on a scheduled Job only this script exists locally, so the module
    # comes from wherever it was imported and the card template from the repo.
    ops.append(CommitOperationAdd("build_local.py", str(Path(__file__).resolve())))
    ops.append(CommitOperationAdd("agent_usage.py", au.__file__))

    card_tpl = HERE / "card" / "README.md"
    if not card_tpl.exists():
        from huggingface_hub import hf_hub_download

        card_tpl = Path(hf_hub_download(PUSH_REPO, "card/README.md", repo_type="dataset"))
    ops.append(CommitOperationAdd("card/README.md", str(card_tpl)))

    # Card template → README.md with the freshness stamp filled in.
    card = (
        card_tpl.read_text()
        .replace("{{LATEST_MONTH}}", latest)
        .replace("{{GENERATED}}", str(date.today()))
    )
    ops.append(CommitOperationAdd("README.md", card.encode()))

    commit = api.create_commit(
        repo_id=PUSH_REPO,
        repo_type="dataset",
        operations=ops,
        commit_message=f"Update derived agent-usage shares (latest month: {latest})",
    )
    print(f"\n✓ pushed {len(ops)} files → {commit.commit_url}")
else:
    print("\n(no AGENT_USAGE_PUSH_REPO set — local build only, nothing pushed)")