claude-code/build_from_raw.py

"""Build TinyGuide dataset from RAW Claude Code transcripts (not skeletons).

Raw transcripts keep full tool inputs AND results — tracebacks, grep output,
test pass/fail, error text — which the skeleton dropped. That lets every rule
fire with real signal. One {prompt, completion} row per tool call.

Usage:
  python build_from_raw.py [glob ...]
  (default sources: ~/.claude/projects/**/*.jsonl and /tmp/cc*/**/*.jsonl)
"""
import json, re, sys, glob, os, collections, random
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from label_rules import choose_label, HINTS
from format_prompt import format_prompt

random.seed(0)
ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "data"

TEST_RE = re.compile(r"\b(pytest|npm test|pnpm test|yarn test|cargo test|go test|bun test|unittest|jest|vitest|tox)\b")
PY_TB_RE = re.compile(r'File "([^"]+\.\w+)", line \d+')
PYTEST_PATH_RE = re.compile(r"\b([\w./-]+\.\w+):\d+\b")
TEST_FAIL_RE = re.compile(r"\b(FAILED|failed|AssertionError|Error|Traceback|\d+ failed|exit code [1-9])", re.I)
TEST_PASS_RE = re.compile(r"\b(\d+ passed|all tests passed|PASSED|\bOK\b|0 failed)\b")
ENV_RE = re.compile(r"command not found|ModuleNotFoundError|No module named|is not recognized|ENOENT|cannot find module|: not found", re.I)
NOTREAD_RE = re.compile(r"not been read yet|Read it first", re.I)
SYMBOL_RE = re.compile(r"(?:NameError|AttributeError|ImportError|cannot import name).*?['\"`](\w+)['\"`]")


def result_text(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return " ".join(x.get("text", "") for x in content if isinstance(x, dict))
    return str(content or "")


def iter_pairs(path):
    """Yield (tool_name, input_dict, result_text, is_error) in order."""
    pending = {}
    try:
        rows = [json.loads(l) for l in open(path) if l.strip()]
    except Exception:
        return
    goal = ""
    for r in rows:
        msg = r.get("message", {}) or {}
        content = msg.get("content")
        if r.get("type") == "user" and isinstance(content, str) and not goal:
            goal = content.strip()
        if not isinstance(content, list):
            continue
        for c in content:
            if not isinstance(c, dict):
                continue
            if c.get("type") == "text" and not goal and r.get("type") == "user":
                goal = c.get("text", "").strip()
            if c.get("type") == "tool_use":
                pending[c.get("id")] = (c.get("name"), c.get("input", {}) or {})
            if c.get("type") == "tool_result":
                tid = c.get("tool_use_id")
                if tid in pending:
                    name, inp = pending.pop(tid)
                    yield goal, name, inp, result_text(c.get("content")), bool(c.get("is_error"))


def first_path(text):
    m = PY_TB_RE.search(text) or PYTEST_PATH_RE.search(text)
    return m.group(1) if m else None


def session_rows(path):
    state = {
        "observed_files": [], "edited_files": [], "traceback_paths": [],
        "command_history": [], "last_test_status": "unknown",
        "code_changed_since_last_test": False, "traceback_symbols": [],
        "last_failure_type": None, "last_error_signature": None,
        "previous_error_signature": None, "last_search_result_class": None,
        "last_test_failed_same_error_after_edit": False,
    }
    traj, rows, goal_final = [], [], "Continue the task."
    for goal, name, inp, res, is_err in iter_pairs(path):
        goal_final = goal or goal_final
        prop = {"name": name, "args": inp}
        # NOTE: no back-to-back cooldown here. Cooldown is a RUNTIME hook
        # concern; baking it into labels gives identical signal patterns two
        # different labels (VERIFY vs NO_HINT) -> ambiguous supervision ->
        # model collapses to NO_HINT on real prompts. Label the true rule.
        key = choose_label(state, prop)
        prompt = format_prompt(goal_final, traj, state, prop)
        rows.append({"key": key, "prompt": prompt, "completion": HINTS[key]})
        _advance(state, traj, name, inp, res, is_err)
    return rows


def _advance(state, traj, name, inp, res, is_err):
    fp = inp.get("file_path")
    cmd = inp.get("command", "") or inp.get("pattern", "")
    short = (fp or cmd or json.dumps(inp))[:60]
    traj.append({"tool": name, "arg": short,
                 "result": ("ERR " + res[:70]) if is_err else (res[:70] or "ok")})
    if len(traj) > 14:
        del traj[0]

    if name == "Read" and fp and not is_err:
        state["observed_files"].append(fp)
    if name in {"Edit", "Write", "MultiEdit", "NotebookEdit"} and fp:
        if NOTREAD_RE.search(res):
            return  # failed edit; nothing changed, state already triggered the hint
        state["edited_files"].append(fp)
        if fp not in state["observed_files"]:
            state["observed_files"].append(fp)
        state["code_changed_since_last_test"] = True
    if name in {"Grep", "Glob"}:
        n = len(re.findall(r"\n", res))
        state["last_search_result_class"] = ("search_no_results" if not res.strip()
                                             else "search_many_results" if n > 30 else "search_few")
    if name == "Bash":
        state["command_history"].append(cmd)
        sig = (res.strip().splitlines() or [""])[-1][:80]
        # failure type reflects only THIS command's result (not sticky)
        state["last_failure_type"] = "missing_package" if ENV_RE.search(res) else None
        if TEST_RE.search(cmd):
            if TEST_FAIL_RE.search(res) and not TEST_PASS_RE.search(res):
                state["last_test_status"] = "failed"
                tb = first_path(res)
                if tb:
                    state["traceback_paths"].append(tb)
                sym = SYMBOL_RE.search(res)
                if sym:
                    state["traceback_symbols"].append(sym.group(1))
                state["last_test_failed_same_error_after_edit"] = (sig == state.get("last_error_signature"))
            else:
                state["last_test_status"] = "passed"
                state["last_failure_type"] = None
            state["code_changed_since_last_test"] = False
        if is_err or TEST_FAIL_RE.search(res):
            state["previous_error_signature"] = state.get("last_error_signature")
            state["last_error_signature"] = sig


def main():
    pats = sys.argv[1:] or [
        os.path.expanduser("~/.claude/projects/**/*.jsonl"),
        "/tmp/cc*/**/*.jsonl", "/tmp/cc*/*.jsonl",
        "/tmp/mimo/**/*.jsonl",
    ]
    files = sorted({f for p in pats for f in glob.glob(p, recursive=True)})
    print(f"transcripts found: {len(files)}")

    sessions = []
    for f in files:
        rs = session_rows(f)
        if rs:
            sessions.append((Path(f).stem, rs))
    print(f"sessions with tool calls: {len(sessions)}")

    random.shuffle(sessions)
    split = int(len(sessions) * 0.9)

    def emit(sess, path, balance):
        allrows = [r for _sid, rs in sess for r in rs]
        # drop consecutive duplicate hints within nothing -> dedup globally per session done below
        hints = [r for r in allrows if r["key"] != "NO_HINT"]
        nohint = [r for r in allrows if r["key"] == "NO_HINT"]
        random.shuffle(hints); random.shuffle(nohint)
        if balance and hints:
            # cap any single hint class to 35% of hints to avoid one rule dominating
            cap = max(50, int(len(hints) * 0.35))
            seen = collections.Counter(); kept = []
            for r in hints:
                if seen[r["key"]] < cap:
                    kept.append(r); seen[r["key"]] += 1
            hints = kept
            # target ~75% NO_HINT
            keep_no = min(len(nohint), int(len(hints) / 0.25 * 0.75))
            nohint = nohint[:keep_no]
        out = hints + nohint
        random.shuffle(out)
        with open(path, "w") as fh:
            for r in out:
                fh.write(json.dumps({"prompt": r["prompt"], "completion": r["completion"]}) + "\n")
        return out

    OUT.mkdir(exist_ok=True)
    tr = emit(sessions[:split], OUT / "train.jsonl", balance=True)
    va = emit(sessions[split:], OUT / "valid.jsonl", balance=True)

    # dump ALL rows (unbalanced, with key), session-split preserved, for build_scaled.py
    def dump_all(sess, path):
        with open(path, "w") as fh:
            for sid, rs in sess:
                for r in rs:
                    fh.write(json.dumps({"key": r["key"], "prompt": r["prompt"],
                                         "completion": r["completion"],
                                         "source": "real", "session_id": sid}) + "\n")
    dump_all(sessions[:split], OUT / "all_train.jsonl")
    dump_all(sessions[split:], OUT / "all_valid.jsonl")

    def dist(rows):
        c = collections.Counter("NO_HINT" if r["completion"] == "NO_HINT" else "HINT" for r in rows)
        return dict(c)
    print(f"train rows: {len(tr)} {dist(tr)}")
    print(f"valid rows: {len(va)} {dist(va)}")
    spread = collections.Counter(r["completion"][:34] for _sid, rs in sessions for r in rs)
    print("label spread (all sessions, pre-balance):")
    for k, n in spread.most_common():
        print(f"  {n:6d}  {k}")


if __name__ == "__main__":
    main()