training/build_cli_decisions.py

#!/usr/bin/env python3
"""Pose CLI Command Risk as typed decisions for the d1 scripts in this folder.

Writes cli-data/decisions/{train,validation,test}.parquet with the columns `common.load_split` reads: id,
task_slug (the template family), state ({shell, command}), questions (17, identical per row) and gold.
Train and validation keep one shell per distinct command and at most --train-cap / --val-cap commands per
family, because rows within a family are near copies. Test keeps every row, so scores compare with the
TF-IDF baseline in cli-command-risk/baseline.py.

    uv run --no-project --with pyarrow python build_cli_decisions.py
"""
import argparse
import collections
import hashlib
import json
import re
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

HERE = Path(__file__).resolve().parent
SOURCE = HERE.parent / "cli-command-risk"
OUT = HERE / "cli-data" / "decisions"
LEVELS = ["safe", "low", "moderate", "high", "critical"]
SHELL_SUFFIX = re.compile(r"-(?:bash|zsh|sh|fish|dash)-[0-9a-f]+$")


def questions(taxonomy, cats):
    defs = {c["name"]: c["definition"] for c in taxonomy["risk_categories"]}
    level_defs = {lv["name"]: lv["definition"] for lv in taxonomy["risk_levels"]}
    out = {
        "risk_level": {"type": "score", "instructions": "How much harm can running this shell command as written cause?",
                       "criteria": [f"{lv.capitalize()}. {level_defs[lv]}" for lv in LEVELS]},
        "risky": {"type": "noul", "instructions": "Is this shell command risky to run as written?",
                  "criteria": {"true": "Running it can destroy data, change permissions or privileges, expose secrets, "
                                       "or disrupt the system.",
                               "false": "It is read-only or routine, with no plausible path to harm."}},
    }
    for c in cats:
        name = c.replace("_", " ")
        out[c] = {"type": "noul", "instructions": f"Does this command carry {name} risk?",
                  "criteria": {"true": defs[c], "false": f"Nothing in the command matches {name}."}}
    out["reversible"] = {"type": "noul", "instructions": "Can an average user undo the effect of this command without backups?",
                         "criteria": {"true": "The effect can be undone without backups, or there is no effect to undo.",
                                      "false": "Undoing the effect needs backups or is not possible."}}
    out["requires_root"] = {"type": "noul", "instructions": "Does the harm from this command need root or sudo?",
                            "criteria": {"true": "The command needs root or sudo to do its harm.",
                                         "false": "An ordinary user can run it with the same effect."}}
    return out


def family(row):
    if row["derived_from"]:
        return row["derived_from"]
    if row["source"] == "open_source":
        return "os:" + row["command"]
    return SHELL_SUFFIX.sub("", row["id"])


def rank(text):
    return hashlib.sha256(text.encode()).hexdigest()


def thin(rows, cap):
    """One shell per distinct command, then the first `cap` commands per family in hash order."""
    by_family = collections.defaultdict(dict)
    for r in rows:
        best = by_family[family(r)].get(r["command"])
        if best is None or rank(r["id"]) < rank(best["id"]):
            by_family[family(r)][r["command"]] = r
    kept = []
    for fam in sorted(by_family):
        kept += sorted(by_family[fam].values(), key=lambda r: rank(r["command"]))[:cap]
    return kept


def gold(row, cats):
    noul = lambda value: {"label": "true" if value else "false"}  # noqa: E731
    out = {"risk_level": {"label": str(LEVELS.index(row["risk_level"]))},
           "risky": noul(row["risk_level"] != "safe")}
    out.update({c: noul(c in row["risk_categories"]) for c in cats})
    out["reversible"] = noul(row["reversible"])
    out["requires_root"] = noul(row["requires_root"])
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train-cap", type=int, default=40)
    ap.add_argument("--val-cap", type=int, default=8)
    ap.add_argument("--full-train", action="store_true", help="keep every train row, all shells, no family cap")
    ap.add_argument("--out", type=Path, default=OUT.parent)
    args = ap.parse_args()
    out = args.out / "decisions"
    taxonomy = json.loads((SOURCE / "taxonomy.json").read_text())
    data = {s: pq.read_table(SOURCE / "data" / f"{s}.parquet").to_pylist() for s in ("train", "validation", "test")}
    present = {c for r in data["train"] for c in r["risk_categories"]}
    cats = [c["name"] for c in taxonomy["risk_categories"] if c["name"] != "safe_baseline" and c["name"] in present]
    qs = json.dumps(questions(taxonomy, cats))
    picked = {"train": data["train"] if args.full_train else thin(data["train"], args.train_cap),
              "validation": thin(data["validation"], args.val_cap), "test": data["test"]}
    out.mkdir(parents=True, exist_ok=True)
    summary = {}
    for split, rows in picked.items():
        table = pa.table({
            "id": [r["id"] for r in rows],
            "task_slug": [family(r) for r in rows],
            "state": [json.dumps({"shell": r["shell"], "command": r["command"]}) for r in rows],
            "questions": [qs] * len(rows),
            "gold": [json.dumps(gold(r, cats)) for r in rows],
        })
        pq.write_table(table, out / f"{split}.parquet")
        summary[split] = {"cases": len(rows), "families": len({family(r) for r in rows}),
                          "levels": dict(collections.Counter(r["risk_level"] for r in rows))}
    summary["questions"] = len(json.loads(qs))
    (out.parent / "summary.json").write_text(json.dumps(summary, indent=1) + "\n")
    print(json.dumps(summary, indent=1))


if __name__ == "__main__":
    main()