scripts/audit_sft_messages.py

#!/usr/bin/env python3
"""Audit SFT message JSONL for training safety."""

from __future__ import annotations

import json
import re
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from build_sft_messages import build_assistant_content
from dataset_quality_utils import (
    asks_for_code,
    asks_hamburger,
    asks_nav_links_open,
    asks_single_file_html,
    has_code,
    has_doctype_or_html,
    has_external_assets,
    has_inline_style_script,
    prompt_overlap_ratio,
)

ROOT = Path(__file__).resolve().parent.parent
AUDIT_DIR = ROOT / "data" / "audit"
REPORT_PATH = AUDIT_DIR / "sft_quality_report.md"
WEAK_SFT_PATH = AUDIT_DIR / "weak_sft_messages.jsonl"


def pick_sft_path() -> Path:
    clean = ROOT / "data" / "train" / "mythos_sft_messages_clean.jsonl"
    default = ROOT / "data" / "train" / "mythos_sft_messages.jsonl"
    return clean if clean.exists() else default


def audit_sft_entry(user: str, assistant: str) -> tuple[list[str], bool]:
    issues: list[str] = []
    hard_fail = False
    a = str(assistant or "")
    u = str(user or "")
    lower_a = a.lower()

    if asks_for_code(u) and not has_code(a):
        issues.append("code_request_no_code")
        hard_fail = True

    if asks_for_code(u) and a.startswith("Diagnosis:"):
        issues.append("code_request_starts_with_diagnosis")
        hard_fail = True

    if prompt_overlap_ratio(u, a) > 0.4:
        issues.append("repeats_user_prompt")
        hard_fail = True

    if "additional targeted edits" in lower_a:
        issues.append("additional_targeted_edits")

    if asks_single_file_html(u):
        if has_external_assets(a):
            issues.append("single_file_external_assets")
            hard_fail = True
        if not has_inline_style_script(a) and has_doctype_or_html(a):
            issues.append("single_file_missing_inline_style_script")
            hard_fail = True

    if asks_nav_links_open(u) or (asks_hamburger(u) and "nav-links" in u.lower()):
        shows_before_after = any(k in lower_a for k in ("broken", "wrong target", "wrong:"))
        has_fixed_toggle = bool(re.search(r"nav-links[^;\n]*classlist\.toggle\(['\"]open", a, re.I))
        if re.search(r'\.hamburger[^;\n]*classlist\.toggle\(["\']open', a, re.I):
            if shows_before_after and has_fixed_toggle:
                pass  # intentional broken vs fixed teaching example
            else:
                issues.append("toggles_open_on_hamburger")
                if not shows_before_after:
                    hard_fail = True
        if not re.search(r'classlist\.toggle\(["\']open["\']', a, re.I):
            issues.append("missing_nav_links_toggle")
            if "show" in u.lower() and "correct" in u.lower() and has_fixed_toggle:
                pass
            elif not shows_before_after:
                hard_fail = True

    if "short test checklist" in lower_a and a.count("-") < 2 and a.count("\n") < 4:
        issues.append("checklist_heading_without_items")

    if a.startswith("Diagnosis:") and "Implementation:" in a:
        impl = a.split("Implementation:", 1)[1].split("Verification:", 1)[0].strip()
        if not impl or (len(impl) < 40 and not has_code(impl)):
            issues.append("diagnosis_without_implementation")

    if "Verification:" in a:
        ver = a.split("Verification:", 1)[1].strip()
        if len(ver) < 15 or ver.lower() in {"verify it works", "confirm the fix holds."}:
            issues.append("vague_verification")

    return issues, hard_fail


def main() -> int:
    AUDIT_DIR.mkdir(parents=True, exist_ok=True)
    sft_path = pick_sft_path()
    if not sft_path.exists():
        print(f"No SFT file at {sft_path}", file=sys.stderr)
        return 1

    entries = []
    weak = []
    issue_counter: Counter = Counter()
    hard_fails = 0

    with sft_path.open("r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            if not line.strip():
                continue
            row = json.loads(line)
            user = row["messages"][1]["content"]
            assistant = row["messages"][2]["content"]
            issues, hard = audit_sft_entry(user, assistant)
            for i in issues:
                issue_counter[i] += 1
            if hard:
                hard_fails += 1
            if issues:
                weak.append({
                    "line_number": line_num,
                    "hard_fail": hard,
                    "issues": issues,
                    "user_prompt": user,
                    "assistant_preview": assistant[:800],
                    "row": row,
                })
            entries.append(row)

    with WEAK_SFT_PATH.open("w", encoding="utf-8") as f:
        for w in weak:
            f.write(json.dumps(w, ensure_ascii=False) + "\n")

    with_code = sum(1 for e in entries if has_code(e["messages"][2]["content"]))
    starts_diagnosis = sum(1 for e in entries if e["messages"][2]["content"].startswith("Diagnosis:"))
    starts_code = sum(1 for e in entries if e["messages"][2]["content"].startswith("Here is the complete code:"))

    lines = [
        "# SFT Quality Report",
        "",
        f"File: `{sft_path}`",
        f"Total SFT rows: **{len(entries)}**",
        f"Assistant with code: **{with_code}**",
        f"Starts with Diagnosis: **{starts_diagnosis}**",
        f"Starts with 'Here is the complete code:': **{starts_code}**",
        f"Weak SFT rows: **{len(weak)}**",
        f"Hard fails: **{hard_fails}**",
        "",
        "## Issue counts",
        "",
    ]
    for issue, count in issue_counter.most_common():
        lines.append(f"- `{issue}`: {count}")

    lines.extend(["", "## Sample weak SFT (first 10)", ""])
    for w in weak[:10]:
        lines.append(f"### line {w['line_number']} hard_fail={w['hard_fail']}")
        lines.append(f"- issues: {', '.join(w['issues'])}")
        lines.append(f"- user: {w['user_prompt'][:200]}")
        lines.append("")

    safe = hard_fails == 0 and (with_code / max(len(entries), 1)) >= 0.5
    lines.extend([
        "",
        "## Training safety",
        "",
        f"**{'SAFE' if safe else 'NOT SAFE'}** to train on this file.",
    ])

    REPORT_PATH.write_text("\n".join(lines), encoding="utf-8")
    print(f"Audited {len(entries)} SFT rows, weak={len(weak)}, hard_fails={hard_fails}")
    print(f"Wrote {REPORT_PATH}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())