diff --git "a/traceforge/cli.py" "b/traceforge/cli.py" new file mode 100644--- /dev/null +++ "b/traceforge/cli.py" @@ -0,0 +1,2757 @@ +"""traceforge CLI — read, benchmark, improve, validate, export. + +The pipeline is an evaluation instrument: it rewrites a model's real traces +into ideal-world references, and `benchmark` scores the model by how little +that rewrite had to change. + +Layout under repo: + data/traces/raw/ normalized model traces (the benchmark's input) + data/traces/improved/ ideal-world references (the benchmark's other half) + data/triage/ per-message labels (keep/drop/rephrase/...) + data/benchmark/ report.json, traces.jsonl, leaderboard.jsonl + data/prompts/ improver prompts (audit trail) + data/work/ per-trace improver workdirs (scratch) + data/manifest.jsonl one status row per improved trace + data/dataset/*.jsonl exported JSONL (downstream reuse only) +""" +from __future__ import annotations + +import argparse +import json +import os +import re +import sqlite3 +import sys +import time +import types +from typing import Optional + +import traceforge +from traceforge import live +from traceforge.config import Config, repo_root_from_pkg +from traceforge.defects import classify +from traceforge.export import export_chat_jsonl +from traceforge.improver.harness import _append_status +from traceforge.readers import ( + ClaudeTraceReader, + CodexTraceReader, + HermesTraceReader, + NemotronTraceReader, +) +from traceforge.trace import Trace +from traceforge.validate import validate_structural + + +def _cfg(args) -> Config: + return Config.load(getattr(args, "config", None), repo_root=repo_root_from_pkg()) + + +SOURCES = ("hermes", "claude", "codex", "nemotron") + + +def _parse_sources(args) -> list[str]: + raw = getattr(args, "source", None) or "all" + if raw == "all": + return list(SOURCES) + parts = [s.strip().lower() for s in raw.split(",") if s.strip()] + bad = [s for s in parts if s not in SOURCES] + if bad: + raise SystemExit(f"error: unknown source(s): {bad} (known: {', '.join(SOURCES)})") + return parts + + +class MultiSourceReader: + """Unified view over every installed agent store. + + load_trace prefers the raw cache (what `pull` wrote) so improve is + reproducible and independent of live stores; unknown ids fall back to + the live readers — hermes first (cheap sqlite lookup), then claude + (JSONL index scan). + """ + + def __init__(self, cfg: Config): + self.cfg = cfg + self._hermes: HermesTraceReader | None = None + self._claude: ClaudeTraceReader | None = None + self._codex: CodexTraceReader | None = None + self._nemotron: NemotronTraceReader | None = None + self._child_ids: set[str] | None = None + + @property + def hermes(self) -> HermesTraceReader | None: + if self._hermes is None: + dbs = [self.cfg.state_db] + list(self.cfg.hermes_state_dbs) + if any(os.path.exists(d) for d in dbs): + self._hermes = HermesTraceReader(dbs) + return self._hermes + + @property + def claude(self) -> ClaudeTraceReader | None: + if self._claude is None: + dirs = [os.path.expanduser("~/.claude/projects")] + list( + self.cfg.claude_projects_dirs + ) + if any(os.path.isdir(d) for d in dirs): + self._claude = ClaudeTraceReader( + dirs, include_sidechains=self.cfg.claude_sidechains + ) + return self._claude + + @property + def codex(self) -> CodexTraceReader | None: + if self._codex is None: + # live store ALWAYS first (owner rule: "live store is always + # read when present"); extra roots add seed archives + roots = [os.path.expanduser("~/.codex/sessions")] + for r in self.cfg.codex_roots: + if r not in roots: + roots.append(r) + if any(os.path.isdir(r) for r in roots): + self._codex = CodexTraceReader(roots) + return self._codex + + @property + def nemotron(self) -> NemotronTraceReader | None: + """Published SFT corpora (scripts/fetch_nemotron.py output).""" + if self._nemotron is None: + roots = [r for r in self.cfg.nemotron_roots if os.path.isdir(r)] + if roots: + self._nemotron = NemotronTraceReader(roots) + return self._nemotron + + def reader_for(self, source: str): + if source == "hermes": + r = self.hermes + if r is None: + raise FileNotFoundError( + f"Hermes state.db not found at {self.cfg.state_db}" + ) + return r + if source == "claude": + r = self.claude + if r is None: + raise FileNotFoundError( + "Claude projects dir not found: ~/.claude/projects" + ) + return r + if source == "codex": + r = self.codex + if r is None: + raise FileNotFoundError( + f"Codex sessions dir(s) not found: {self.cfg.codex_roots}" + ) + return r + if source == "nemotron": + r = self.nemotron + if r is None: + raise FileNotFoundError( + "Nemotron seed dir not found: " + f"{self.cfg.nemotron_roots or 'seed-datasets/nemotron'} " + "(run scripts/fetch_nemotron.py)" + ) + return r + raise KeyError(source) + + # ------------------------------------------------------------ listing + + def list_sessions(self, *, sources=None, **kw) -> list[dict]: + sources = sources or SOURCES + rows: list[dict] = [] + guard_kw = dict(kw) + excl_kw = guard_kw.pop("exclude_improver_children", None) + for src in sources: + try: + r = self.reader_for(src) + except FileNotFoundError as e: + print(f"-- source '{src}' unavailable: {e}") + continue + kwd = dict(guard_kw) + # Self-reference guard at the SQL level: HermesTraceReader + # drops rows whose cwd sits inside the improver workdir — + # mis-tagged improver children otherwise look like real + # candidates and improve the model on its own output. + if src == "hermes" and (excl_kw is None or excl_kw): + kwd.setdefault("exclude_cwd", self.cfg.improver_workdir) + # seed-archive discovery can take a while on first run; + # heartbeat so long parses don't look hung (only readers + # that accept status_cb — codex's parallel discovery) + import inspect + + if "status_cb" in inspect.signature( + r.list_sessions + ).parameters: + def _status(msg, _src=src): + print(f"-- {_src} {msg}", flush=True) + kwd.setdefault("status_cb", _status) + got = r.list_sessions(**kwd) + for s in got: + s["_reader"] = src # serving reader (s['source'] = provenance) + rows.extend(got) + rows.sort(key=lambda s: s["started_at"] or 0.0, reverse=True) + if kw.get("limit") is not None: + rows = rows[: kw["limit"]] + return rows + + def sources_of(self, sid: str) -> list[str]: + """Which stores can serve this session id (raw cache excluded).""" + out = [] + if self.hermes is not None: + try: + self.hermes.load_trace(sid) + out.append("hermes") + except KeyError: + pass + if self.claude is not None and sid in self.claude._index(): + out.append("claude") + if self.codex is not None and sid in self.codex._index(): + out.append("codex") + if self.nemotron is not None: + try: + self.nemotron.load_trace(sid) + out.append("nemotron") + except KeyError: + pass + return out + + # ------------------------------------------------------------- loading + + def load_trace(self, sid: str) -> Trace: + cached = os.path.join(self.cfg.raw_dir, f"{sid}.json") + if os.path.exists(cached): + return Trace.load(cached) + # live stores: hermes first (cheap sqlite lookup) + if self.hermes is not None: + try: + return self.hermes.load_trace(sid) + except KeyError: + pass + if self.claude is not None: + try: + return self.claude.load_trace(sid) + except KeyError: + pass + if self.codex is not None: + try: + return self.codex.load_trace(sid) + except KeyError: + pass + if self.nemotron is not None: + try: + return self.nemotron.load_trace(sid) + except KeyError: + pass + raise KeyError(f"session not found in any source: {sid}") + + # --------------------------------------------------- self-reference guard + + def improver_child_ids(self) -> set[str]: + """Hermes session ids that traceforge's own pipeline spawned. + + Every improver stage (writer, digest, map, verify, review) is a + Hermes one-shot running with cwd inside cfg.improver_workdir. + Improving such a trace trains the model on its own output — an + endless self-referential loop — so these ids are non-candidates + everywhere: listing, pulling, and improving. + """ + if self._child_ids is not None: + return self._child_ids + ids: set[str] = set() + wd = (self.cfg.improver_workdir or "").rstrip("/\\") + like = wd.replace("\\", "\\\\").replace("%", "\\%").replace( + "_", "\\_" + ) + "/%" + for db in [self.cfg.state_db] + list(self.cfg.hermes_state_dbs): + if not os.path.exists(db): + continue + try: + conn = sqlite3.connect( + f"file:{db}?mode=ro", uri=True + ) + try: + ids |= { + row[0] + for row in conn.execute( + "SELECT id FROM sessions WHERE cwd LIKE ?" + " ESCAPE '\\'", + (like,), + ) + } + finally: + conn.close() + except Exception: + pass + self._child_ids = ids + return ids + + def is_improver_child(self, sid: str, trace: Optional[Trace] = None) -> bool: + """True when sid is one of our improver's own Hermes sessions + (cwd-based, the only signal that survives mis-tagged sources).""" + wd = (self.cfg.improver_workdir or "").rstrip("/\\") + if wd and trace is not None and (trace.cwd or "").startswith(wd + os.sep): + return True + return sid in self.improver_child_ids() + + +def _print_sessions(sessions) -> None: + for s in sessions: + started = time.strftime("%Y-%m-%d %H:%M", time.gmtime(s["started_at"])) if s["started_at"] else "?" + title = (s["title"] or "")[:52] + print( + f"{s['session_id']:38} {s.get('source', '?'):7} msgs={s['message_count']:4} " + f"tools={s['tool_call_count']:4} {started} {title}" + ) + print(f"-- {len(sessions)} candidate traces") + + +def _pull_sessions(cfg, r, sessions) -> int: + """Pull the listed candidate sessions into data/traces/raw/. + + Over-long traces (rendered size > cfg.split_max_chars) are split + into trainable segment traces #p1..N at assistant-turn + boundaries (call→result groups never torn apart; the task row is + prepended to every segment) — one giant 500k-token training + sequence becomes several ≤~120k-token ones, each flowing through + the full quality pipeline on its own. + """ + os.makedirs(cfg.raw_dir, exist_ok=True) + from traceforge.progress import Progress + from traceforge.splitter import needs_split, save_segments, split_trace + + n = 0 + # A split session writes #p1..N.json and NEVER .json, so the + # plain existence check below re-pulled and re-split every split + # session on every pass — measured 2026-10-23: ~2,000 sessions + # (including 480KB monsters) re-parsed and re-split each pass, ~15 min + # of dead time before the improve stage could start. Index the segment + # bases once instead of scanning the raw dir per session. + existing = set(os.listdir(cfg.raw_dir)) + seg_bases = {f.split("#p")[0] for f in existing if "#p" in f} + with Progress(len(sessions), desc="pull:") as prog: + for s in sessions: + sid = s["session_id"] + out = os.path.join(cfg.raw_dir, f"{sid}.json") + if not cfg.overwrite and (f"{sid}.json" in existing + or sid in seg_bases): + prog.write(f"skip {sid} (already pulled)") + prog.update() + continue + t = r.reader_for(s["_reader"]).load_trace(sid) + trigger = cfg.split_trigger_chars or cfg.split_max_chars + if cfg.split_long and needs_split(t, split_max_chars=trigger): + segs = split_trace( + t, + split_max_chars=cfg.split_max_chars, + min_segment_chars=cfg.min_segment_chars, + max_segments=cfg.max_segments_per_trace, + ) + sids = save_segments(segs, cfg.raw_dir, s["session_id"]) + n += len(sids) + prog.write( + f"split {s['source']:7} {s['session_id']} -> " + f"{len(segs)} segments ({', '.join(x.split('#')[-1] for x in sids)})" + ) + else: + t.save(out) + n += 1 + prog.write(f"pulled {s['source']:7} {s['session_id']} -> data/traces/raw/") + prog.update() + print(f"-- {n} traces pulled ({len(sessions)} sessions)") + return 0 + + +def _pick(args, cfg, name: str): + """CLI flag wins, else the config value. + + The argparse defaults used to be 10/120/2, which shadowed + config.example.yaml at every call site — so `min_messages`, + `max_messages` and `min_tool_calls` in the YAML were dead settings. + Defaults are now None and the config is the fallback. + """ + v = getattr(args, name, None) + return v if v is not None else getattr(cfg, name) + + +def cmd_list(args) -> int: + cfg = _cfg(args) + r = MultiSourceReader(cfg) + sessions = r.list_sessions( + sources=_parse_sources(args), + min_messages=_pick(args, cfg, "min_messages"), + max_messages=_pick(args, cfg, "max_messages"), + min_tool_calls=_pick(args, cfg, "min_tool_calls"), + limit=args.limit, + ) + _print_sessions(sessions) + return 0 + + +def cmd_pull(args) -> int: + cfg = _cfg(args) + r = MultiSourceReader(cfg) + sessions = r.list_sessions( + sources=_parse_sources(args), + min_messages=_pick(args, cfg, "min_messages"), + max_messages=_pick(args, cfg, "max_messages"), + min_tool_calls=_pick(args, cfg, "min_tool_calls"), + limit=args.limit, + ) + return _pull_sessions(cfg, r, sessions) + + +def _crash_result(sid: str, exc: Exception): + """Manifest-shaped status row for a worker crash (never loses a batch).""" + return { + "session_id": sid, "ok": False, + "errors": [f"worker crashed: {exc}"], + "hermes_session_id": None, "attempts": 0, "duration_s": 0.0, + } + + +_MANIFEST_LOCK = __import__("threading").Lock() + + +def _write_manifest_row(cfg, row: dict) -> None: + """Append ONE outcome row the moment the attempt finishes. + + The batch used to buffer every row and flush at the END of the run + (see _improve_ids), so a batch that ran for hours — or was killed — + left the manifest with stage rows but almost no outcomes: measured + 2026-10-23, 4628 manifest rows and only 54 carried ok/errors/attempts + against 3633 attempted traces. Absence of a failure row therefore did + NOT mean no failure, and the funnel's failure census was blind. + """ + try: + os.makedirs(os.path.dirname(cfg.manifest_path), exist_ok=True) + with _MANIFEST_LOCK: + with open(cfg.manifest_path, "a", encoding="utf-8") as f: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + except OSError: + pass # bookkeeping must never kill the batch + + +def _improve_one(cfg, r, sid, max_attempts, skip_existing, prog, liv=None): + """Improve one trace; returns (sid, ImproverResult). Thread-safe.""" + if skip_existing and os.path.exists( + os.path.join(cfg.improved_dir, f"{sid}.json") + ): + prog.write(f"SKIP {sid} (already improved)") + return sid, None + if liv is not None: + liv.registry.begin(sid) + t = r.load_trace(sid) + from traceforge.improver import improve_trace + + def cb(s: str) -> None: + if liv is not None: + liv.registry.set_stage(sid, s) + else: + prog.set_postfix(f"{sid[:12]}… {s}") + + try: + if cfg.splice: + # windowed direct-HTTP engine (no agent runtime) + from traceforge.improver import ImproverResult + from traceforge.splice import improve_trace_splice + + sres = improve_trace_splice( + cfg, t, status_cb=cb, max_attempts=max_attempts, + ) + res = ImproverResult( + session_id=sid, ok=sres["ok"], improved=sres["improved"], + reply=sres.get("route", ""), errors=sres["errors"], + attempts=sres["attempts"], + duration_s=sres["duration_s"], + ) + if res.ok and res.improved is not None \ + and sres.get("route") != "regate": + out = os.path.join(cfg.improved_dir, f"{sid}.json") + res.improved.save(out) + prog.write( + f"IMPROVED {sid} ({res.attempts} attempts, " + f"{res.duration_s:.0f}s, {res.reply})" + ) + elif res.ok: + prog.write( + f"IMPROVED {sid} (regate, {res.duration_s:.0f}s)" + ) + else: + prog.write(f"FAILED {sid}: {res.errors}") + _write_manifest_row(cfg, res.to_dict()) + if liv is not None: + liv.registry.end(sid) + return sid, res + res = improve_trace( + cfg, t, max_attempts=max_attempts, status_cb=cb + ) + except Exception as e: + if liv is not None: + liv.registry.end(sid) + _write_manifest_row(cfg, _crash_result(sid, e)) + raise + if liv is not None: + liv.registry.end(sid) + toks = "" + if liv is not None: + d = liv.per_trace(sid) + if d and (d.get("in") or d.get("out")): + toks = f" · {live.fmt_tok(d['in'])} in · {live.fmt_tok(d['out'])} out" + if d.get("rate"): + toks += f" · {d['rate']:.1f} tok/s" + if res.ok and res.improved is not None: + out = os.path.join(cfg.improved_dir, f"{sid}.json") + res.improved.save(out) + prog.write( + f"IMPROVED {sid} ({res.attempts} attempts, {res.duration_s:.0f}s" + f"{toks})" + ) + else: + prog.write(f"FAILED {sid}{toks}: {res.errors}") + _write_manifest_row(cfg, res.to_dict()) + return sid, res + + +def cmd_improve(args) -> int: + cfg = _cfg(args) + if getattr(args, "no_triage", False): + cfg.triage = False + r = MultiSourceReader(cfg) + ids = args.session_ids or [] + if args.from_raw: + ids = [ + f[:-5] + for f in os.listdir(cfg.raw_dir) + if f.endswith(".json") + ] + # Self-reference guard: improver children never get improved, even + # when a pre-guard run cached them into data/traces/raw/. + child_ids = r.improver_child_ids() + if child_ids: + drop = [i for i in ids if i in child_ids] + if drop: + ids = [i for i in ids if i not in child_ids] + print( + f"-- skipped {len(drop)} traceforge improver session(s) " + f"in input (self-referential data): " + + ", ".join(drop[:5]) + (" …" if len(drop) > 5 else "") + ) + if not ids: + print("no traces to improve — pass session ids, or use --from-raw") + return 2 + if not ids: + print("no traces to improve — pass session ids, or use --from-raw") + return 2 + ids = _filter_improver_children_live(cfg, r, ids) + if not ids: + print("-- no improvable traces left after self-reference guard") + return 0 + return _improve_ids( + cfg, r, ids, + max_attempts=args.max_attempts, + workers=args.workers, + skip_existing=args.skip_existing, + ) + + +def _raw_cwd(path: str) -> Optional[str]: + """Read only the top-level `cwd` from a raw trace file. + + Trace.to_dict writes cwd as the 5th key, so it lands in the first few + hundred bytes. Parsing the whole file (up to 480KB x 12.7k traces) + just to read that one field cost ~15 min of dead time at the start of + every pass. Returns None when the head does not carry it, so the + caller falls back to a full load. + """ + try: + with open(path, "r", encoding="utf-8") as f: + head = f.read(4096) + except OSError: + return None + m = re.search(r'"cwd"\s*:\s*"((?:[^"\\]|\\.)*)"', head) + if not m: + return None + try: + return json.loads('"' + m.group(1) + '"') + except Exception: + return m.group(1) + + +def _filter_improver_children_live(cfg, r, ids) -> list[str]: + """Drop ids that a currently-running batch (or its state.db rows) + marked as improver children since the reader snapshot was taken, and + drop raw-cached traces whose own Trace.cwd sits in the workdir.""" + wd = (cfg.improver_workdir or "").rstrip("/\\") + child_ids = r.improver_child_ids() + cand = [sid for sid in ids if sid not in child_ids] + if not wd: + return cand + # Resolve cwd from the raw head in parallel. This runs at the start of + # every pass over the whole pending set; 36k sequential 4KB opens cost + # 25+ min on Windows with real-time AV scanning, and threads cut it to + # about a minute. Only the head read is threaded: the live-store + # readers share sqlite handles and are not safe to call concurrently. + from concurrent.futures import ThreadPoolExecutor + + def _head_cwd(sid): + return _raw_cwd(os.path.join(cfg.raw_dir, f"{sid}.json")) + + with ThreadPoolExecutor(max_workers=32) as ex: + heads = list(ex.map(_head_cwd, cand)) + out = [] + for sid, cwd in zip(cand, heads): + if cwd is None: + try: + cwd = r.load_trace(sid).cwd or "" + except Exception: + out.append(sid) # unknown store rows: let downstream report + continue + if cwd.startswith(wd + os.sep) or cwd.rstrip("/\\") == wd: + continue + out.append(sid) + return out + + +def cmd_regate(args) -> int: + """Adjudicate existing workdir artifacts with the windowed judge.""" + from concurrent.futures import ThreadPoolExecutor, as_completed + + cfg = _cfg(args) + if not cfg.splice: + print("-- regate requires the splice engine (splice: true)") + return 2 + work_dir = cfg.improver_workdir + if not os.path.isdir(work_dir): + print("no workdir — nothing to regate") + return 1 + ids = args.session_ids or [ + d for d in os.listdir(work_dir) + if os.path.exists(os.path.join(work_dir, d, "improved.json")) + ] + promoted = {f[:-5] for f in os.listdir(cfg.improved_dir) + if f.endswith(".json")} if os.path.isdir(cfg.improved_dir) else set() + ids = [i for i in ids if i not in promoted] + if not ids: + print("no unpromoted artifacts to regate") + return 0 + from traceforge.progress import Progress + from traceforge.splice import judge_artifact + + ok = 0 + workers = max(1, args.workers) + print(f"-- regating {len(ids)} artifacts with the windowed judge " + f"({workers} workers)") + with Progress(len(ids), desc="regate:") as prog: + if workers == 1 or len(ids) == 1: + for sid in ids: + res = judge_artifact(cfg, sid) + if res.get("ok"): + ok += 1 + prog.write(f"PROMOTED {sid} ({res['duration_s']:.0f}s" + + (", cached" if res.get("cached") else "") + + ")") + else: + prog.write(f"FAIL {sid}: {res['errors']}") + prog.update() + else: + with ThreadPoolExecutor(max_workers=workers) as ex: + futs = {ex.submit(judge_artifact, cfg, sid): sid + for sid in ids} + for fut in as_completed(futs): + sid = futs[fut] + try: + res = fut.result() + except Exception as e: + prog.write(f"FAIL {sid}: worker crashed: {e}") + res = {"ok": False} + if res.get("ok"): + ok += 1 + prog.write( + f"PROMOTED {sid} ({res['duration_s']:.0f}s" + + (", cached" if res.get("cached") else "") + + ")") + else: + prog.write(f"FAIL {sid}: {res.get('errors')}") + prog.update() + print(f"-- regate done: {ok}/{len(ids)} promoted") + return 0 if ok or not ids else 1 + + +def cmd_crossexam(args) -> int: + """Cross-examine promoted traces with a second judge family.""" + from concurrent.futures import ThreadPoolExecutor, as_completed + + cfg = _cfg(args) + if not cfg.splice: + print("-- crossexam requires the splice engine (splice: true)") + return 2 + model = args.model + if not model: + print("-- cross-exam needs an explicit --model from a DIFFERENT " + "family than the improver (e.g. gemma4-31b-selfhosted when " + "the improver is qwen-fable-selfhosted)") + return 2 + imp_dir = cfg.improved_dir + if not os.path.isdir(imp_dir): + print("no promoted corpus") + return 1 + ids = args.session_ids or sorted( + f[:-5] for f in os.listdir(imp_dir) if f.endswith(".json")) + if not ids: + print("no promoted traces") + return 0 + from traceforge.progress import Progress + from traceforge.splice import cross_examine + + disagree = [] + n_pass = 0 + workers = max(1, args.workers) + print(f"-- cross-examining {len(ids)} promoted traces with " + f"{model} ({workers} workers)") + with Progress(len(ids), desc="crossexam:") as prog: + if workers == 1 or len(ids) == 1: + for sid in ids: + res = cross_examine(cfg, sid, model) + if res.get("ok"): + n_pass += 1 + else: + disagree.append(res) + prog.update() + else: + with ThreadPoolExecutor(max_workers=workers) as ex: + futs = {ex.submit(cross_examine, cfg, sid, model) + : sid for sid in ids} + for fut in as_completed(futs): + sid = futs[fut] + try: + res = fut.result() + except Exception as e: + res = {"ok": False, "sid": sid, + "violations": [{"detail": + f"worker crashed: {e}"}]} + if res.get("ok"): + n_pass += 1 + prog.write(f"CONFIRM {sid}") + else: + disagree.append(res) + vs = res.get("violations") or [{}] + prog.write( + f"DISAGREE {sid}: " + + "; ".join(str(v.get("detail", ""))[:90] + for v in vs[:2])) + prog.update() + print(f"-- crossexam done: {n_pass}/{len(ids)} confirmed, " + f"{len(disagree)} disagreements") + if disagree: + out = os.path.join(cfg.data_dir, "crossexam_disagreements.json") + os.makedirs(cfg.data_dir, exist_ok=True) + with open(out, "w", encoding="utf-8") as f: + json.dump([{ + "sid": d.get("sid"), "verdict": d.get("verdict"), + "answer_quality": d.get("answer_quality"), + "violations": d.get("violations"), + } for d in disagree], f, indent=1, ensure_ascii=False) + print(f"-- disagreements written to {out}") + return 0 + + +def cmd_triage(args) -> int: + """Label traces via the cheap LLM classifier; prints the route.""" + from traceforge import triage as triage_mod + + cfg = _cfg(args) + r = MultiSourceReader(cfg) + ids = args.session_ids or [] + if args.from_raw: + ids = [ + f[:-5] + for f in os.listdir(cfg.raw_dir) + if f.endswith(".json") + ] + if not ids: + print("no traces to triage — pass session ids, or use --from-raw") + return 2 + counts = {"splice": 0, "agent": 0, "reject": 0} + routes = {} + from traceforge.progress import Progress + + with Progress(len(ids), desc="triage:") as prog: + for sid in ids: + try: + t = r.load_trace(sid) + except Exception as e: + prog.write(f"ERROR {sid}: {e}") + prog.update() + continue + try: + rec = triage_mod.triage_trace(cfg, t, force=args.force) + except RuntimeError as e: + print(f"error: {e}", file=sys.stderr) + return 2 + if rec is None: + prog.write(f"FAILED {sid}: all windows failed") + prog.update() + continue + route = rec["route"] + routes[sid] = route + counts[route] = counts.get(route, 0) + 1 + nl = len(rec["labels"]) + prog.write( + f"{route.upper():7} {sid} ({nl} labels, " + f"{rec['windows_ok']}/{rec['n_windows']} windows)" + ) + prog.update() + print( + f"-- {len(routes)} triaged: " + + ", ".join(f"{v} {k}" for k, v in counts.items() if v) + ) + return 0 if routes else 1 + + +def _improve_ids(cfg, r, ids, *, max_attempts, workers, skip_existing) -> int: + """Shared improve driver for cmd_improve and cmd_run. + + Queue order (2026-09-28, measured funnel): smallest raw trace first. + Promotion yield is strongly inversely correlated with raw size + (corpus: ≤60KB ≈ 27% promoted, ≥300KB ≈ 1%) and a first-queue bias + previously sent 633–733KB monsters ahead of 40KB sure-things — a + thin improved/ dir early in a batch is NOT a gate regression; + checking raw sizes is. Reverse to put the high-yield cohort through + the funnel before the expensive low-yield tail. + """ + def _raw_size(sid: str) -> int: + try: + return os.path.getsize(os.path.join(cfg.raw_dir, f"{sid}.json")) + except OSError: + return 1 << 30 # unknown ids (store, not cache) sink to the tail + + # Queue order: NEVER-ATTEMPTED first, then smallest raw first, with + # multi-turn as a TIEBREAKER inside a size (2026-10-24). + # + # Size alone (the 2026-09-28 rule) is right about yield but wrong about + # exploration: pending = raw minus promoted, so every previously-failed + # trace is retried on every pass, and the small ones come first. Measured + # 2026-10-22: 4530 pending, 1360 already attempted (mostly unfixable — + # subagent tails with no evidence, degenerate 1-6KB traces), 3170 never + # attempted and sitting behind them in the 180KB+ bands. The batch spent + # hours re-churning known failures at ~0 promotions while 1825 traces of + # 300KB+ had never been looked at once. A workdir is the attempt marker + # (it persists after a failure), so bucket on it and keep size order + # inside each bucket: explore the pool, then retry the failures. + # An infra failure is not an attempt (2026-10-26, suite R-92). + # `judge infra error:` / llm_unavailable is 807 of 2,839 coded failures + # (28.5%) and it arrives in endpoint-outage BURSTS, not as a property of + # the trace: measured 10-26, 191 sessions in the 02h, 397 in the 07h, + # 122 in the 08h, median raw 20 KB. Because a workdir is the attempt + # marker, every one of those sank to the retry tail behind 33,791 + # pending raws, and 0 of the 807 had been promoted 4 h later — the + # cheapest cohort in the pool, parked for weeks. Rank it FIRST (0), ahead + # of never-attempted (1) and genuinely-failed (2): it is the cheapest + # cohort in the pool (median 21 KB), its triage cache is already warm so + # the retry skips the pre-pass, and it was lost to an outage rather than + # to trace quality. An outage must cost the pass, not the session. + _att_cache: dict[str, int] = {} + + def _attempted(sid: str) -> int: + if sid in _att_cache: + return _att_cache[sid] + wd = os.path.join(cfg.improver_workdir, sid) + if not os.path.isdir(wd): + rank = 1 + else: + rank = 2 + try: + with open(os.path.join(wd, "failure.json"), + encoding="utf-8") as fh: + rec = json.load(fh) + code = rec.get("defect") or classify( + stage=rec.get("stage"), + errors=rec.get("errors") or [], + violations=rec.get("violations") or []) + # An outage is not a property of the trace, so an infra + # failure is retried ahead of never-attempted work — but + # only in a class that can pay. Measured 2026-10-26: the + # infra cohort was 161 nemotron-ka (46% yield) and 496 + # nemotron-tc (14.9%), and rank 0 put all 657 ahead of + # 1,946 never-attempted real sessions (85.6%). The tc + # half alone is ~2 h of 8-worker time for ~74 traces. + # Keep the bonus where the class yields; let the rest + # queue by class like everything else. + if code == "D-INF-01" and _class_rank(sid) <= 2: + rank = 0 + except (OSError, ValueError): + pass + _att_cache[sid] = rank + return rank + + # Class yield, measured 2026-10-26 on the live pool. Size order alone + # put the WORST class at the front, because the nemotron tool_calling + # subset is the smallest thing in the pool: 4,737 raws, 2,383 attempted, + # 355 promoted = 14.9%, against 85.6% for a real captured session + # (1,432 attempted, 1,226 promoted). The queue was spending its front on + # the 14.9% class while the 85.6% class waited behind it. + # + # Rank is the measured yield, best first. Size stays as the tiebreaker + # INSIDE a class, where it is still the right signal. + # A split segment is a slice of a session, so its yield depends on the + # PARENT: a session that already produced one promotable segment is + # material the gates accept, and its other segments are worth trying. + # Measured 2026-10-26: 548 parents produced a promotion, and 2,649 of + # 18,943 pending segments (14%) sit under one of them. The other 86% + # are slices of sessions that have never yielded anything, at ~2.9%. + # Ranking the productive parents first spends the queue on the 14% that + # can pay instead of draining 18,943 segments in file order. + _good_parents = { + sid.split("#p")[0] for sid in + (os.listdir(cfg.improved_dir) if os.path.isdir(cfg.improved_dir) else []) + if "#p" in sid + } + + def _class_rank(sid: str) -> int: + if sid.startswith("nemotron-"): + return 3 if "agentic-tc" in sid else 2 + if "#sa-" in sid: + return 1 # subagent: ~40% of attempted + if "#p" in sid: + # 4 = segment of a session that has already promoted something, + # 5 = segment of a session that never has + return 4 if sid.split("#p")[0] in _good_parents else 5 + return 0 # real captured session: 85.6% + + # Multi-turn was the FIRST key (2026-10-22, owner ask: more multi-turn + # traces). Measured 2026-10-24 it was inverting the whole queue: the + # multi-turn cohort is the long-session cohort, so 52% of pending + # multi-turn traces are 480KB+ — the first 50 ids of a random pending + # sample were 86% 480KB+ (measured prom/wh 0.00-0.36) while the 60-180KB + # single-turn cohort (3.0-3.3 prom/wh) sat at the tail. Multi-turn also + # promotes WORSE per hour than single-turn: 13.4% at 620s mean (0.79 + # prom/wh) vs 27.3% at 296s (3.32 prom/wh) over the last 24h. So it is + # now a tiebreaker inside a size, not a queue-wide override: multi-turn + # still wins among equally cheap traces, but it can no longer jump the + # 480KB+ monsters ahead of the high-yield small cohort. + _mt_cache: dict[str, int] = {} + # The key is a TIEBREAKER inside a size bucket, but computing it means + # reading every raw in full: 12,897 files / 4.52 GB, measured ~10 min on + # a cold cache (and it blocks the heartbeat, so the pass looks hung). + # Persist it once; a re-pulled raw keeps its old value, which is fine + # for a tiebreaker. + _mt_path = os.path.join(os.path.dirname(cfg.manifest_path), + ".multi_turn.json") + try: + with open(_mt_path, encoding="utf-8") as fh: + _mt_cache = {str(k): int(v) for k, v in json.load(fh).items()} + except (OSError, ValueError, AttributeError): + _mt_cache = {} + _mt_dirty = [False] + + def _multi_turn(sid: str) -> int: + if sid not in _mt_cache: + try: + with open(os.path.join(cfg.raw_dir, f"{sid}.json"), + encoding="utf-8", errors="replace") as fh: + n = fh.read().count('"role": "user"') + except OSError: + n = 0 + _mt_cache[sid] = 0 if n >= 2 else 1 + _mt_dirty[0] = True + return _mt_cache[sid] + + ids = sorted(ids, key=lambda s: (_attempted(s), _class_rank(s), + _raw_size(s), _multi_turn(s))) + if _mt_dirty[0]: + try: + tmp = _mt_path + ".tmp" + with open(tmp, "w", encoding="utf-8") as fh: + json.dump(_mt_cache, fh) + os.replace(tmp, _mt_path) + except OSError: + pass + os.makedirs(cfg.improved_dir, exist_ok=True) + manifest_path = cfg.manifest_path + workers = max(1, workers or 1) + + # Cheap triage pass first: label every message keep/drop/rephrase/ + # fact_check/summarize via direct structured-generation LLM calls. + # splice-route traces skip the improver entirely (spliced downstream + # when the splicer lands); reject-route traces drop out here. + # + # PARALLEL (2026-10-23). This used to be a serial loop over every + # pending trace — one trace at a time, each parallelising only its own + # windows — so the pre-pass was a hard gate in front of the first + # promotion. Measured: 21 triage traces/min with ~5,400 pending after + # the pool doubled = 4-6 h of silent log before any improve work. The + # per-trace work is self-contained (own cache file, own endpoint + # handle, own window pool), so the outer loop parallelises safely; + # triage_prepass_workers bounds it so the endpoint is not swamped. + # The pre-pass is only needed for the AGENT path. improve_trace_splice + # triages each trace itself (stage 1, cached), so pre-triaging the whole + # pool just moves the same cost in front of the first promotion — and + # with the pool doubled it became the dominant cost: measured + # 2026-10-23, 5,224 untriaged traces x ~17 windows = ~89k label calls + # = ~16 h of silent log before any improve work. With splice on, triage + # happens inside the workers and promotions start immediately. + triage_routes: dict[str, str] = {} + if cfg.triage and not cfg.splice: + from concurrent.futures import ThreadPoolExecutor + + from traceforge import triage as triage_mod + + def _triage_one(sid: str) -> tuple[str, str, Optional[str]]: + """(sid, route, fatal_error) — never raises.""" + try: + t = r.load_trace(sid) + except Exception: + return sid, "agent", None # unresolvable: default route + try: + rec = triage_mod.triage_trace(cfg, t) + except RuntimeError as e: + return sid, "agent", str(e) + if rec is None: + return sid, "agent", None + return sid, rec["route"], None + + outer = max(1, int(getattr(cfg, "triage_prepass_workers", 4) or 4)) + fatal: Optional[str] = None + with ThreadPoolExecutor(max_workers=outer) as ex: + for sid, route, err in ex.map(_triage_one, ids): + triage_routes[sid] = route + if err and fatal is None: + fatal = err + if fatal: + print(f"-- triage unavailable ({fatal}); improving without labels") + triage_routes = {sid: "agent" for sid in ids} + n_splice = sum(1 for v in triage_routes.values() if v == "splice") + n_reject = sum(1 for v in triage_routes.values() if v == "reject") + n_agent = sum(1 for v in triage_routes.values() if v == "agent") + print( + f"-- triage: {n_agent} agent · {n_splice} splice · " + f"{n_reject} reject" + ) + if n_reject: + from traceforge.improver.harness import _append_status + + for sid in ids: + if triage_routes.get(sid) == "reject": + _append_status(cfg, sid, { + "stage": "triage", "verdict": "REJECT", + "labels": 0, + }) + ids = [sid for sid in ids if triage_routes.get(sid) != "reject"] + if not ids: + print("-- every trace rejected at triage; nothing to improve") + return 0 + + if workers > 1 and len(ids) > 1: + print(f"-- improving {len(ids)} traces, {workers} concurrent workers") + statuses = [] + from traceforge.progress import Progress + + with Progress(len(ids), desc="improve:") as prog: + # live telemetry: poll child-session token counters in state.db; + # every run gets it (single trace shows its tok/s subline too) + liv = live.LiveStats.setup(cfg, prog) + if liv is not None: + liv.start() + try: + if workers == 1 or len(ids) == 1: + for sid in ids: + _, res = _improve_one( + cfg, r, sid, max_attempts, skip_existing, + prog, liv, + ) + if res is not None: + statuses.append(res) + prog.update() + else: + from concurrent.futures import ThreadPoolExecutor, as_completed + + # load traces serially (SQLite conn per call), improve in parallel + with ThreadPoolExecutor(max_workers=workers) as ex: + futures = { + ex.submit( + _improve_one, cfg, r, sid, max_attempts, + skip_existing, prog, liv, + ): sid + for sid in ids + } + for fut in as_completed(futures): + sid = futures[fut] + try: + _, res = fut.result() + except Exception as e: # never lose a whole batch + prog.write(f"FAILED {sid}: worker crashed: {e}") + res = _crash_result(sid, e) + if res is not None: + statuses.append(res) + prog.update() + finally: + if liv is not None: + liv.stop() + if liv is not None: + print(liv.summary()) + # Outcome rows are appended per trace by _improve_one (and by the + # crash handler below) the moment each attempt finishes — a batch that + # runs for hours or is killed must not lose its failure evidence. + from traceforge.improver.harness import ImproverResult + + rows = [ + res.to_dict() if isinstance(res, ImproverResult) else res + for res in statuses + ] + skipped = len(ids) - len(statuses) + if not rows: + # nothing was attempted: all skipped (idempotent resume) = done OK; + # nothing selectable at all = user error + if skipped: + print(f"-- {skipped} traces skipped (all already improved)") + return 0 + return 2 + ok = sum(1 for s in rows if s["ok"]) + print(f"-- {ok}/{len(rows)} traces improved" + ( + f", {skipped} skipped" if skipped else "")) + return 0 if ok == len(rows) else 1 + + +def cmd_audit(args) -> int: + """Mechanical, LLM-free audit of the promoted corpus.""" + from traceforge.audit import main as audit_main + + cfg = _cfg(args) + target = args.files[0] if args.files else cfg.improved_dir + return audit_main(target, verbose=args.verbose) + + +def cmd_validate(args) -> int: + cfg = _cfg(args) + raw_dir = cfg.raw_dir + from traceforge.grounding import validate_grounding + + bad = 0 + files = args.files or [ + os.path.join(cfg.improved_dir, f) + for f in sorted(os.listdir(cfg.improved_dir)) + if f.endswith(".json") + ] + for fp in files: + t = Trace.load(fp) + errs = validate_structural(t) + # grounding needs the source trace; skip silently when raw is absent + raw_path = os.path.join(raw_dir, f"{t.parent_session_id}.json") + if os.path.exists(raw_path): + errs += validate_grounding(t, Trace.load(raw_path)) + if errs: + bad += 1 + print(f"INVALID {os.path.basename(fp)}:") + for e in errs[:10]: + print(f" - {e}") + else: + print(f"ok {os.path.basename(fp)}") + return 1 if bad else 0 + + +def cmd_verify(args) -> int: + """Re-run external fact verification on an existing improved trace.""" + cfg = _cfg(args) + from traceforge.improver.verify import run_verifier + from traceforge.progress import Progress + + files = args.files or [ + os.path.join(cfg.improved_dir, f) + for f in sorted(os.listdir(cfg.improved_dir)) + if f.endswith(".json") + ] + bad = 0 + with Progress(len(files), desc="verify:") as prog: + liv = live.LiveStats.setup(cfg, prog) + if liv is not None: + liv.start() + try: + for fp in files: + sid = os.path.basename(fp)[:-5].replace("ideal-", "") + workdir = os.path.join(cfg.improver_workdir, sid) + os.makedirs(workdir, exist_ok=True) + if liv is not None: + liv.registry.begin(sid) + liv.registry.set_stage(sid, "verify") + else: + prog.set_postfix(os.path.basename(fp)[:24]) + res = run_verifier(cfg, fp, workdir) + if liv is not None: + liv.registry.end(sid) + tag = os.path.basename(fp) + if res.ok(): + prog.write( + f"PASS {tag} ({res.claims_checked} external" + f" claims checked, {res.duration_s:.0f}s)" + ) + elif res.verdict == "FAIL": + bad += 1 + prog.write(f"FAIL {tag}") + for u in res.unsupported[:10]: + prog.write(f" - {u.get('claim', '?')}: {u.get('why', '?')}") + else: + bad += 1 + prog.write(f"ERROR {tag}: {res.error}") + _append_status(cfg, sid, { + "stage": "verify", "verdict": res.verdict, + "claims_checked": res.claims_checked, + }) + prog.update() + finally: + if liv is not None: + liv.stop() + if liv is not None: + print(liv.summary()) + return 1 if bad else 0 + + + + +def _pct(vals: list[int], q: float) -> int: + """Nearest-rank percentile (q in 0..1) over a non-empty list.""" + if not vals: + return 0 + s = sorted(vals) + i = min(len(s) - 1, max(0, int(round(q * (len(s) - 1))))) + return s[i] + + +def _dist(vals: list[int]) -> dict: + """mean/median/p90/max for a per-trace metric.""" + if not vals: + return {"mean": 0, "median": 0, "p90": 0, "max": 0} + return { + "mean": round(sum(vals) / len(vals), 1), + "median": _pct(vals, 0.5), + "p90": _pct(vals, 0.9), + "max": max(vals), + } + + +def _ts_summary(traces) -> dict: + """Time-sensitivity census for the dataset card (never raises).""" + try: + from traceforge.timesensitivity import summary + return summary(traces) + except Exception as e: # a card must never block an export + return {"error": str(e)} + + +def _write_dataset_card(traces, out: str, val_frac: float, + dedup_removed: int = 0) -> None: + """Measured stats the training run should know: distribution biases, + shape, and the guarantees the funnel enforces.""" + import re as _re + from collections import Counter as _C + from traceforge.export import _agent_role + langs = _C() + for t in traces: + task = next((m.content or "" for m in t.messages + if m.role == "user"), "") + if _re.search(r"\b(der|die|das|und|nicht|bitte|f\u00fcr|mit|ist)\b", + task, _re.I): + langs["de"] += 1 + else: + langs["en"] += 1 + # orchestrator/subagent balance (MASS: train on both conversation + # roles, balanced — a blind sample underweights the subagent half) + agent_roles = _C(r for r in (_agent_role(t) for t in traces) if r) + # task/topic balance: the trainer needs to know which buckets are thin + # before it weights them (office-docs/creative are near-empty by design + # of the seed pool, not by classifier failure) + tasks = _C(t.task for t in traces if t.task) + + # ---- shape: what a training example actually looks like ------------- + # Counted per trace so the trainer can size its context window and + # batch budget from the real distribution instead of the mean. + role_tot = _C() + chars_by_role = _C() + msgs_per_trace: list[int] = [] + calls_per_trace: list[int] = [] + users_per_trace: list[int] = [] + chars_per_trace: list[int] = [] + tool_names = _C() + # Verified recovery label (docs/DATA_QUALITY.md §4.1): exactly one + # environment-caused failure followed by exactly one corrective step. + # The earlier heuristic ("an error-looking tool row and a later assistant + # turn") counted grep exit-1 and error text quoted inside source code as + # failures — it reported 32% where the real figure is 1.8%. + from traceforge.export import recovery_label + recovery_traces = 0 + for t in traces: + n_calls = 0 + n_users = 0 + n_chars = 0 + for m in t.messages: + role_tot[m.role] += 1 + chars_by_role[m.role] += len(m.content or "") + if m.role == "user": + n_users += 1 + n_chars += len(m.content or "") + for c in (m.tool_calls or []): + n_calls += 1 + tool_names[c.name] += 1 + n_chars += len(c.arguments or "") + chars_by_role["assistant"] += len(c.arguments or "") + if recovery_label(t): + recovery_traces += 1 + msgs_per_trace.append(len(t.messages)) + calls_per_trace.append(n_calls) + users_per_trace.append(n_users) + chars_per_trace.append(n_chars) + + # Token estimate: no tokenizer is a dependency of this repo, so use the + # standard chars/4 prose proxy and SAY SO. Tool output is code/JSON/logs, + # which tokenizes denser than prose, so treat this as a lower bound. + # The budget the owner trains at is the model's full 128k window, so that + # is the headline number; the smaller budgets stay as reference points. + est_tokens = [c // 4 for c in chars_per_trace] + fits = {str(n): sum(1 for v in est_tokens if v <= n) + for n in (4096, 8192, 16384, 65536, 131072)} + total_chars = sum(chars_by_role.values()) or 1 + + card = { + "traces": len(traces), + # docs/DATA_QUALITY.md rule 9: publish the dedup removal rate. + "dedup": { + "removed": dedup_removed, + "pct": round(100.0 * dedup_removed + / max(1, len(traces) + dedup_removed), 1), + "key": "(task, tool-names + sha1(args))", + }, + "languages": dict(langs), + "agent_roles": dict(agent_roles), + "tasks": dict(tasks.most_common()), + "messages": { + "total": sum(msgs_per_trace), + "by_role": dict(role_tot.most_common()), + "per_trace": _dist(msgs_per_trace), + }, + "tool_calls": { + "total": sum(calls_per_trace), + "per_trace": _dist(calls_per_trace), + "traces_with_none": sum(1 for v in calls_per_trace if v == 0), + "distinct_tools": len(tool_names), + "top_tools": dict(tool_names.most_common(15)), + }, + "user_turns": { + "per_trace": _dist(users_per_trace), + "single_turn_traces": sum(1 for v in users_per_trace if v <= 1), + "multi_turn_traces": sum(1 for v in users_per_trace if v > 1), + }, + "tokens": { + "method": "chars/4 (prose proxy). Calibrated against a real " + "tokenizer on this corpus: chars/4 runs ~13% low " + "(median 21.4k est vs 24.5k actual), so scale by ~1.15 " + "for a central estimate; tool output is code/JSON and " + "tokenizes denser than prose.", + "budget": 131072, + "chars_total": sum(chars_per_trace), + "est_total": sum(est_tokens), + "per_trace": _dist(est_tokens), + "fits_budget": { + k: {"traces": v, "pct": round(100.0 * v / max(1, len(traces)), 1)} + for k, v in fits.items() + }, + }, + "split": {"val_frac": val_frac} if val_frac else None, + # The masking contract is part of the data, not a trainer detail: a + # trainer that computes loss on tool results trains on the + # environment's text instead of the policy's (docs/DATA_QUALITY.md §4.3). + "supervision": { + "contract": "loss on assistant spans only; tool/user/system " + "spans are masked context", + "chars_by_role": dict(chars_by_role), + "assistant_share_pct": round( + 100.0 * chars_by_role["assistant"] / total_chars, 1), + "tool_share_pct": round( + 100.0 * chars_by_role["tool"] / total_chars, 1), + "recovery": { + "traces": recovery_traces, + "pct": round(100.0 * recovery_traces / max(1, len(traces)), 1), + "cap_pct": 15.0, + "over_cap": recovery_traces > 0.15 * max(1, len(traces)), + "method": "verified label: exactly one environment-caused " + "failure (flaky endpoint, missing dependency, " + "exhausted resource) followed by exactly one " + "corrective step; model-caused errors and retry " + "loops disqualify (docs/DATA_QUALITY.md §4.1)", + }, + }, + # docs/DATA_QUALITY.md §4.3: filtering to clean paths skews difficulty, + # so the distribution is reported rather than assumed. + "difficulty": { + "tool_calls": _dist(calls_per_trace), + "messages": _dist(msgs_per_trace), + "est_tokens": _dist(est_tokens), + "bands": { + "calls_lt_10": sum(1 for v in calls_per_trace if v < 10), + "calls_10_24": sum(1 for v in calls_per_trace if 10 <= v < 25), + "calls_25_49": sum(1 for v in calls_per_trace if 25 <= v < 50), + "calls_ge_50": sum(1 for v in calls_per_trace if v >= 50), + }, + }, + # Time sensitivity: a temporal fact answered from weights decays. + # Measured on the improved Trace, which carries no `tools` array, so + # a trace whose harness recorded no schemas cannot be distinguished + # from one with no retrieval tool — the count is an UPPER bound. + # `traceforge timesensitivity --dataset train.jsonl` measures the + # export (which does carry `tools`) and is the number to quote. + "time_sensitivity": _ts_summary(traces), + "notes": [ + "identity/refusal burns gated (traceforge/identity.py)", + "clip markers, template leaks, permission-ask finals gated", + "residual-defect estimate ~5% (selfhosted cross-exam, triaged)", + "reasoning channel dropped per message (keep_reasoning=False)", + ], + } + card_path = os.path.join(os.path.dirname(out), "dataset_card.json") + with open(card_path, "w", encoding="utf-8") as f: + json.dump(card, f, indent=2, ensure_ascii=False) + print(f" dataset card -> {card_path}") + # one-line shape summary so the numbers are visible without opening JSON + print(f" shape: {card['messages']['total']} msgs " + f"({card['messages']['per_trace']['median']} median/trace), " + f"{card['tool_calls']['total']} tool calls " + f"({card['tool_calls']['per_trace']['median']} median/trace), " + f"~{card['tokens']['est_total'] / 1e6:.1f}M est tokens " + f"({card['tokens']['per_trace']['median']} median/trace, " + f"{card['tokens']['fits_budget']['131072']['pct']}% fit 128k, " + f"max {card['tokens']['per_trace']['max']})") + + +def _tasks_path(cfg) -> str: + return os.path.join(cfg.data_dir, "tasks.jsonl") + + +def _load_task_labels(cfg) -> dict[str, dict]: + """sid -> {task, method, confidence} from data/tasks.jsonl.""" + path = _tasks_path(cfg) + out: dict[str, dict] = {} + if not os.path.exists(path): + return out + with open(path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + r = json.loads(line) + except json.JSONDecodeError: + continue + if r.get("session_id") and r.get("task"): + out[r["session_id"]] = r + return out + + +def _classify_files(cfg, files, *, use_llm: bool = True, force: bool = False, + workers: int = 8, quiet: bool = False) -> dict[str, dict]: + """Label `files` and merge into data/tasks.jsonl. Returns new rows. + + Shared by `cmd_classify` and the generate pass, so a batch run labels + its own promotions without a manual step. + """ + from concurrent.futures import ThreadPoolExecutor, as_completed + + from traceforge.classify import classify_trace + + existing = _load_task_labels(cfg) + todo = [] + for p in files: + t = Trace.load(p) + if t.session_id in existing and not force: + continue + todo.append((p, t)) + if not todo: + return {} + + if not quiet: + print(f"classifying {len(todo)} traces " + f"({'llm+heuristic' if use_llm else 'heuristic only'}, " + f"{workers} workers)") + + def _work(item): + p, t = item + return p, t, classify_trace(t, cfg, use_llm=use_llm) + + rows: dict[str, dict] = {} + done = 0 + with ThreadPoolExecutor(max_workers=max(1, workers)) as ex: + futs = [ex.submit(_work, it) for it in todo] + for fut in as_completed(futs): + p, t, r = fut.result() + rows[t.session_id] = { + "session_id": t.session_id, + "task": r["task"], + "method": r.get("method", "heuristic"), + "confidence": r.get("confidence", 0.0), + "reason": r.get("reason", ""), + "heuristic_task": (r.get("heuristic") or {}).get("task"), + } + done += 1 + if not quiet and done % 50 == 0: + print(f" {done}/{len(todo)}") + + merged = {**existing, **rows} + path = _tasks_path(cfg) + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as f: + for sid in sorted(merged): + f.write(json.dumps(merged[sid], ensure_ascii=False) + "\n") + os.replace(tmp, path) + return rows + + +def cmd_classify(args) -> int: + """Label improved traces with a task bucket. + + Labels land in data/tasks.jsonl (sid -> task/method/confidence) rather + than mutating the improved artifacts: the corpus stays byte-identical, + labels are reviewable in one file, and re-labeling is a rewrite of one + small file instead of 600. + """ + from traceforge.classify import TASKS + + cfg = _cfg(args) + src_dir = cfg.improved_dir + if args.files: + files = list(args.files) + else: + files = [os.path.join(src_dir, f) for f in sorted(os.listdir(src_dir)) + if f.endswith(".json")] + if not files: + print("no improved traces to classify — run `improve` first") + return 2 + + rows = _classify_files( + cfg, files, + use_llm=not getattr(args, "heuristic_only", False), + force=getattr(args, "force", False), + workers=max(1, getattr(args, "workers", 8) or 1), + ) + if not rows: + print(f"all {len(files)} traces already labeled " + f"(use --force to reclassify)") + return 0 + + merged = _load_task_labels(cfg) + import collections + dist = collections.Counter(r["task"] for r in merged.values()) + print(f"\nlabeled {len(rows)} new · {len(merged)} total -> " + f"{_tasks_path(cfg)}") + for task in TASKS: + n = dist.get(task, 0) + if n: + print(f" {task:14s} {n:5d} {n / len(merged):5.1%}") + disagree = sum(1 for r in rows.values() + if r["method"] == "llm" + and r["heuristic_task"] + and r["heuristic_task"] != r["task"]) + print(f" llm/heuristic disagreement: {disagree}/{len(rows)}") + return 0 + + +def _export_by_topic(traces, out: str, cfg, keep_reasoning: bool) -> dict[str, int]: + """Write one train..jsonl per bucket next to `out`. + + Labels come from data/tasks.jsonl (`traceforge classify`); traces with + no label are classified with the free heuristic so the split is always + complete. Returns {task: rows}. + """ + from traceforge.classify import TASKS, classify_heuristic + + labels = _load_task_labels(cfg) + buckets: dict[str, list] = {t: [] for t in TASKS} + for t in traces: + task = (labels.get(t.session_id) or {}).get("task") + if not task: + task = classify_heuristic(t)["task"] + t.task = task + buckets.setdefault(task, []).append(t) + + out_dir = os.path.dirname(os.path.abspath(out)) + stem = os.path.splitext(os.path.basename(out))[0] + counts: dict[str, int] = {} + for task, group in buckets.items(): + if not group: + continue + path = os.path.join(out_dir, f"{stem}.{task}.jsonl") + counts[task] = export_chat_jsonl( + group, path, keep_reasoning=keep_reasoning) + return counts + + +def cmd_export(args) -> int: + cfg = _cfg(args) + src_dir = cfg.improved_dir + # export precondition: the corpus must pass the mechanical audit — + # a corpus with self-talk/clip/unparseable findings must not reach + # training silently (force with --force-audit-fail). + if not getattr(args, "force_audit_fail", False): + from traceforge.audit import main as audit_main + if audit_main(src_dir) != 0: + print("-- export blocked: audit findings above; fix them " + "or pass --force-audit-fail to override") + return 1 + traces = [] + for f in sorted(os.listdir(src_dir)): + if f.endswith(".json"): + traces.append(Trace.load(os.path.join(src_dir, f))) + if not traces: + print("no improved traces to export — run `improve` first") + return 2 + # docs/DATA_QUALITY.md rule 9: dedup at the (task, tool-sequence) level + # before quality filtering, and publish the removal rate. A re-pulled + # session or a re-improved split segment lands here as an exact repeat. + # The sequence carries a sha1 of the call ARGUMENTS (export.dedup_key), + # not just the tool names — name-only collapsed distinct sessions. + from traceforge.export import dedup_traces + traces, dupes = dedup_traces(traces) + if dupes: + print(f"-- dedup: removed {len(dupes)} duplicate trace(s) " + f"({100.0 * len(dupes) / max(1, len(traces) + len(dupes)):.1f}%) " + f"by (task, tool-names + sha1(args))") + # Time sensitivity (opt-in filter): a temporal fact the agent had no + # retrieval tool for is a decaying answer. Training on it teaches the + # model to ASSERT the fact from its weights, which is wrong the moment + # the fact changes. The same fact fetched with a tool is a good target. + if getattr(args, "drop_time_sensitive", False): + from traceforge.timesensitivity import is_publishable + kept = [t for t in traces if is_publishable(t)] + dropped = len(traces) - len(kept) + if dropped: + print(f"-- time sensitivity: dropped {dropped} trace(s) whose " + f"answer is a temporal fact with no retrieval tool") + traces = kept + out = args.out or os.path.join(cfg.data_dir, "dataset", "train.jsonl") + val_frac = max(0.0, min(0.5, getattr(args, "val_frac", 0.0) or 0.0)) + if val_frac: + # deterministic split: stable per-session hash, no randomness — + # re-exports with a grown corpus keep prior assignments stable + import hashlib + train_traces, val_traces = [], [] + for t in traces: + h = int(hashlib.sha256( + (t.parent_session_id or t.session_id or "").encode() + ).hexdigest()[:8], 16) + (val_traces if (h % 1000) < int(val_frac * 1000) + else train_traces).append(t) + val_out = os.path.join(os.path.dirname(out), "val.jsonl") + n_val = export_chat_jsonl( + val_traces, val_out, + keep_reasoning=getattr(args, "keep_reasoning", False)) + n = export_chat_jsonl( + train_traces, out, + keep_reasoning=getattr(args, "keep_reasoning", False)) + print(f"exported {n} train -> {out}, {n_val} val -> {val_out} " + f"(val_frac={val_frac:.2f})") + if getattr(args, "by_topic", False): + counts = _export_by_topic( + train_traces, out, cfg, getattr(args, "keep_reasoning", False)) + print(" by topic (train split):") + for task, c in sorted(counts.items(), key=lambda kv: -kv[1]): + print(f" {task:14s} {c:5d}") + _write_dataset_card(traces, out, val_frac, dedup_removed=len(dupes)) + return 0 + n = export_chat_jsonl( + traces, out, keep_reasoning=getattr(args, "keep_reasoning", False) + ) + print(f"exported {n} traces -> {out}") + if getattr(args, "by_topic", False): + counts = _export_by_topic( + traces, out, cfg, getattr(args, "keep_reasoning", False)) + print(" by topic:") + for task, c in sorted(counts.items(), key=lambda kv: -kv[1]): + print(f" {task:14s} {c:5d} -> " + f"{os.path.basename(out).replace('.jsonl', '')}.{task}.jsonl") + # corpus profile: surface known distribution biases so the trainer + # can weight/curate (measured 2026-10-05: EN-heavy, review/fix skew) + import re as _re + from collections import Counter as _C + langs = _C() + for t in traces: + task = next((m.content or "" for m in t.messages + if m.role == "user"), "") + if _re.search(r"\b(der|die|das|und|nicht|bitte|für|mit|ist)\b", + task, _re.I): + langs["de"] += 1 + else: + langs["en"] += 1 + print(f" corpus profile: {dict(langs)}") + _write_dataset_card(traces, out, val_frac, dedup_removed=len(dupes)) + return 0 + + +def cmd_kto(args) -> int: + """Preference corpus (KTO) — docs/DATA_QUALITY.md §4.4, rule 10. + + Two stages, because KTO needs rejected samples from the policy being + aligned, not relabelled rows from the SFT corpus: + + --prompts extract rollout seeds from the SFT corpus (cleaned prefixes) + --build turn labelled rollouts into length-matched KTO rows + + Stage 2 (sampling + verification) needs a served model and runs between + the two: sample N completions per prompt from the SFT model, verify each + against the trace's own tool evidence, and label it desirable/undesirable. + """ + cfg = _cfg(args) + from traceforge.export import build_kto, build_kto_provisional, extract_prompts + out_dir = os.path.join(cfg.data_dir, "dataset") + os.makedirs(out_dir, exist_ok=True) + + if args.provisional: + src_dir = cfg.improved_dir + traces = [Trace.load(os.path.join(src_dir, f)) + for f in sorted(os.listdir(src_dir)) if f.endswith(".json")] + if not traces: + print("no improved traces — run `improve` first") + return 2 + rows, stats = build_kto_provisional( + traces, + os.path.join(cfg.data_dir, "traces", "raw"), + os.path.join(cfg.data_dir, "triage"), + max_rows=args.max_prompts, + ) + out = args.out or os.path.join(out_dir, "kto_provisional.jsonl") + with open(out, "w", encoding="utf-8") as f: + for r in rows: + f.write(json.dumps(r, ensure_ascii=False) + "\n") + print(f"kto provisional: {stats['rows']} rows -> {out}") + print(f" desirable {stats['desirable']} / undesirable " + f"{stats['undesirable']} (ratio {stats['ratio']})") + print(f" NOT on-policy — {stats['note']}") + return 0 + + if args.build: + rollouts = [] + with open(args.build, encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + rollouts.append(json.loads(line)) + rows, stats = build_kto(rollouts, max_len_ratio=args.max_len_ratio) + out = args.out or os.path.join(out_dir, "kto.jsonl") + with open(out, "w", encoding="utf-8") as f: + for r in rows: + f.write(json.dumps(r, ensure_ascii=False) + "\n") + print(f"kto: {stats['rows']} rows -> {out}") + print(f" desirable {stats['desirable']} / undesirable " + f"{stats['undesirable']} over {stats['prompts']} prompts") + if stats["dropped_length_mismatch"]: + print(f" dropped {stats['dropped_length_mismatch']} negative(s) " + f"for length mismatch (>{args.max_len_ratio}x the shortest " + f"desirable completion)") + if stats["dropped_unpaired"]: + print(f" dropped {stats['dropped_unpaired']} negative(s) with no " + f"desirable sample for the same prompt") + return 0 + + src_dir = cfg.improved_dir + traces = [Trace.load(os.path.join(src_dir, f)) + for f in sorted(os.listdir(src_dir)) if f.endswith(".json")] + if not traces: + print("no improved traces — run `improve` first") + return 2 + prompts = extract_prompts(traces, max_prompts=args.max_prompts, + turns=args.turns, + max_prefix_chars=args.max_prefix_chars) + out = args.out or os.path.join(out_dir, "kto_prompts.jsonl") + with open(out, "w", encoding="utf-8") as f: + for p in prompts: + f.write(json.dumps(p, ensure_ascii=False) + "\n") + print(f"kto prompts: {len(prompts)} rollout seeds -> {out}") + print(" next: sample N completions per prompt from the SFT model, " + "verify each against the trace's tool evidence, label it, then " + "`traceforge kto --build `") + return 0 + + +def cmd_stats(args) -> int: + cfg = _cfg(args) + r = MultiSourceReader(cfg) + out = {"sources": {}} + for src in _parse_sources(args): + try: + out["sources"][src] = r.reader_for(src).stats() + except FileNotFoundError as e: + out["sources"][src] = {"error": str(e)} + print(json.dumps(out, indent=2)) + return 0 + + +def cmd_funnel(args) -> int: + """Census + gate re-run over the current data state. + + Answers "where are traces dropping?" in one command: raw pool vs + workdir stage-reach vs promoted, a sampled re-adjudication of + unpromoted writer artifacts against the live gates (how many are + actually clean = free promotion candidates vs clip-marker vs + real failures), per-size yield, and manifest stage tallies. Read-only. + """ + import random + import collections + + from traceforge.grounding import validate_grounding + from traceforge.improver.harness import ( + _parse_improved, _clip_markers_present, + ) + + cfg = _cfg(args) + raw_dir, work_dir, imp_dir = cfg.raw_dir, cfg.improver_workdir, cfg.improved_dir + raw = {f[:-5] for f in os.listdir(raw_dir) if f.endswith(".json")} \ + if os.path.isdir(raw_dir) else set() + promoted = {f[:-5] for f in os.listdir(imp_dir) if f.endswith(".json")} \ + if os.path.isdir(imp_dir) else set() + + print(f"pool: {len(raw)} raw · {len(promoted)} promoted " + f"({len(promoted) / len(raw):.1%} of pool)" if raw else "no raw pool") + # multi-agent composition (MASS doctrine): orchestrator vs subagent + # halves of the pool — the census shows whether both roles flow. + # (Disjoint: a split segment of a sidechain counts as subagent.) + if raw: + n_sub = sum(1 for sid in raw if "#sa-" in sid) + n_seg = sum(1 for sid in raw if "#p" in sid and "#sa-" not in sid) + comp = f" · {n_sub} subagent" if n_sub else "" + comp += f" · {n_seg} split-segment" if n_seg else "" + if comp: + print(f"composition:{comp}") + never = raw - set(os.listdir(work_dir)) if os.path.isdir(work_dir) else raw + print(f"never attempted: {len(never)}") + + # workdir stage census + stage = collections.Counter() + sizes = {sid: os.path.getsize(os.path.join(raw_dir, f"{sid}.json")) + for sid in raw} + unpromoted_artifacts = [] + if os.path.isdir(work_dir): + for d in os.listdir(work_dir): + wd = os.path.join(work_dir, d) + if not os.path.isdir(wd): + continue + has_imp = os.path.exists(os.path.join(wd, "improved.json")) + if d in promoted: + stage["promoted"] += 1 + continue + if has_imp: + unpromoted_artifacts.append(d) + if any(f.startswith("review") for f in os.listdir(wd)): + stage["died at/after review"] += 1 + elif os.path.exists(os.path.join(wd, "verdict.json")): + stage["died at/after verify"] += 1 + else: + stage["writer artifact, pre-verify"] += 1 + else: + stage["no writer artifact"] += 1 + for k, v in stage.most_common(): + print(f" workdirs: {k}: {v}") + + # failure census: WHY the empty workdirs are empty. The splice engine + # used to persist nothing on failure, so this was unknowable; every + # _fail() now writes data/work//failure.json (stage + violations). + fail_stage = collections.Counter() + fail_pillar = collections.Counter() + fail_err = collections.Counter() + n_fail_rec = 0 + if os.path.isdir(work_dir): + for d in os.listdir(work_dir): + fp = os.path.join(work_dir, d, "failure.json") + if not os.path.exists(fp): + continue + try: + with open(fp, encoding="utf-8") as f: + rec = json.load(f) + except Exception: + continue + n_fail_rec += 1 + fail_stage[str(rec.get("stage"))] += 1 + for v in (rec.get("violations") or []): + fail_pillar[str(v.get("pillar"))] += 1 + for e in (rec.get("errors") or [])[:1]: + head = str(e).split(":")[0][:60] + fail_err[head] += 1 + if n_fail_rec: + print(f"failure records: {n_fail_rec}") + print(" by stage:") + for k, v in fail_stage.most_common(12): + print(f" {k}: {v}") + if fail_pillar: + print(" by cited pillar:") + for k, v in fail_pillar.most_common(10): + print(f" {k}: {v}") + print(" by first error:") + for k, v in fail_err.most_common(10): + print(f" {k}: {v}") + bar_denom = len(raw) + attempt_proxy = bar_denom - len(never) + print(f"attempted ≈ {attempt_proxy} " + f"({attempt_proxy and f'{len(promoted)/max(1,attempt_proxy):.1%} promoted of attempted'})") + + # size-yield buckets + if raw: + buckets: dict[int, list[int]] = {} + for sid in raw: + kb = min(sizes.get(sid, 0), 300_000) // 60_000 # 0..5 + buckets.setdefault(kb, [0, 0]) + buckets[kb][0] += 1 + if sid in promoted: + buckets[kb][1] += 1 + print("yield by raw size:") + for kb in sorted(buckets): + tot, ok = buckets[kb] + lo = kb * 60 + label = f"{lo}-{lo + 60}KB" if kb < 5 else "300KB+" + print(f" {label:>12}: {ok}/{tot}") + + # sampled gate re-adjudication of unpromoted artifacts + n = min(args.sample, len(unpromoted_artifacts)) + if n: + random.seed(args.sample) + picked = random.sample(unpromoted_artifacts, n) + verdicts = collections.Counter() + for sid in picked: + art = os.path.join(work_dir, sid, "improved.json") + t = _parse_improved(art) + if t is None: + verdicts["unparseable artifact"] += 1 + continue + if _clip_markers_present(t): + verdicts["clip markers (repair)"] += 1 + continue + raw_path = os.path.join(raw_dir, f"{sid}.json") + if not os.path.exists(raw_path): + verdicts["raw missing"] += 1 + continue + try: + real = Trace.load(raw_path) + errs = validate_structural(t) or validate_grounding(t, real) + except Exception as e: + verdicts[f"gate crash: {type(e).__name__}"] += 1 + continue + verdicts["CLEAN — free promotion candidate" if not errs + else "gate failure (repair with errors)"] += 1 + scale = len(unpromoted_artifacts) / n + print(f"unpromoted artifacts: {len(unpromoted_artifacts)} " + f"(sampled {n}, scaled ×{scale:.1f}):") + for k, v in verdicts.most_common(): + print(f" {k}: {v} ≈ {int(round(v * scale))}") + + # manifest tallies + if os.path.exists(cfg.manifest_path): + tallies = collections.Counter() + with open(cfg.manifest_path, encoding="utf-8") as f: + for line in f: + try: + r = json.loads(line) + except Exception: + continue + if "stage" in r: + tallies[f"{r.get('stage')}|{r.get('verdict', '?')}"] += 1 + if tallies: + print("manifest stage rows:") + for k, v in tallies.most_common(): + print(f" {k}: {v}") + return 0 + + +def cmd_benchmark(args) -> int: + """Score a model's traces by how little the pipeline had to edit them. + + Read-only over the on-disk corpus: every promoted trace is a + (model trace, ideal rewrite) pair. Fidelity = similarity of the pair; + labels = the triage keep/drop judgement; composite = weighted mean. + The inverse framing is the point: a trace the pipeline can promote + unchanged is a trace the model already got right. + """ + from traceforge.benchmark import ( + DEFAULT_WEIGHTS, append_leaderboard, build_report, load_leaderboard, + load_triage, render_leaderboard, render_report, score_pair, + ) + + cfg = _cfg(args) + raw_dir = args.raw_dir or cfg.raw_dir + imp_dir = args.improved_dir or cfg.improved_dir + triage_dir = args.triage_dir or cfg.triage_dir + if not os.path.isdir(imp_dir): + print(f"no improved traces at {imp_dir} — run `improve` first") + return 2 + + w_sim = max(0.0, args.w_similarity) + w_lab = max(0.0, args.w_labels) + if w_sim + w_lab <= 0: + print("error: --w-similarity and --w-labels cannot both be zero") + return 2 + weights = {"similarity": w_sim, "labels": w_lab} + + # Gate attempts are sparse (only traces whose improve run wrote a + # manifest row carry them); the benchmark reports them when present and + # never imputes them. + attempts_by_sid: dict[str, int] = {} + if os.path.exists(cfg.manifest_path): + with open(cfg.manifest_path, "r", encoding="utf-8") as f: + for line in f: + try: + r = json.loads(line) + except ValueError: + continue + if "attempts" in r and r.get("session_id"): + attempts_by_sid[r["session_id"]] = r["attempts"] + + # Task labels live in data/tasks.jsonl keyed by the improved session id + # (`ideal-`), not on the artifact — the corpus stays byte-identical + # and re-labelling rewrites one small file. + task_by_sid: dict[str, str] = {} + tasks_path = os.path.join(cfg.data_dir, "tasks.jsonl") + if os.path.exists(tasks_path): + with open(tasks_path, "r", encoding="utf-8") as f: + for line in f: + try: + r = json.loads(line) + except ValueError: + continue + if r.get("session_id") and r.get("task"): + task_by_sid[r["session_id"]] = r["task"] + + sources = None if (args.source or "all") == "all" else set(_parse_sources(args)) + files = sorted(f for f in os.listdir(imp_dir) if f.endswith(".json")) + if args.limit: + files = files[: args.limit] + + scores = [] + skipped_no_raw = 0 + skipped_source = 0 + for f in files: + try: + improved = Trace.load(os.path.join(imp_dir, f)) + except (OSError, ValueError, KeyError): + continue + sid = improved.parent_session_id or improved.session_id.replace("ideal-", "") + raw_path = os.path.join(raw_dir, f"{sid}.json") + if not os.path.exists(raw_path): + skipped_no_raw += 1 + continue + try: + raw = Trace.load(raw_path) + except (OSError, ValueError, KeyError): + skipped_no_raw += 1 + continue + if sources is not None and raw.source not in sources: + skipped_source += 1 + continue + improved.task = (improved.task or task_by_sid.get(improved.session_id) + or task_by_sid.get(sid)) + triage = load_triage(os.path.join(triage_dir, f"{sid}.json")) + scores.append(score_pair( + raw, improved, triage=triage, + attempts=attempts_by_sid.get(sid), weights=weights, + )) + + if not scores: + print(f"no scoreable pairs in {imp_dir} (raw dir: {raw_dir})") + return 2 + + report = build_report( + scores, label=args.label, weights=weights, + corpus={ + "raw_pool": (len(os.listdir(raw_dir)) + if os.path.isdir(raw_dir) else 0), + "improved_pool": len(files), + "skipped_no_raw": skipped_no_raw, + "skipped_source": skipped_source, + }, + ) + + out = args.out or os.path.join(cfg.data_dir, "benchmark", "report.json") + os.makedirs(os.path.dirname(os.path.abspath(out)), exist_ok=True) + with open(out, "w", encoding="utf-8") as f: + json.dump(report, f, indent=2, ensure_ascii=False) + traces_path = os.path.join(os.path.dirname(os.path.abspath(out)), + "traces.jsonl") + with open(traces_path, "w", encoding="utf-8") as f: + for s in sorted(scores, key=lambda s: s.score): + f.write(json.dumps(s.to_dict(), ensure_ascii=False) + "\n") + + if args.json: + print(json.dumps(report, indent=2, ensure_ascii=False)) + else: + print(render_report(report)) + print(f"\nreport -> {out}\nper-trace -> {traces_path}") + + if args.label: + lb_path = os.path.join(os.path.dirname(os.path.abspath(out)), + "leaderboard.jsonl") + append_leaderboard(lb_path, report) + print("\n" + render_leaderboard(load_leaderboard(lb_path))) + return 0 + + +def _quarantined_sids(cfg) -> set[str]: + """Session ids parked under data/quarantine/ (all repair classes). + + The quarantine tree is CLASSED — data/quarantine//.json + with class in echo/, rework-echo/, rework-language/, rework-reframe/ + — so a flat listdir of data/quarantine returns ZERO .json files and + the skip set silently came back empty. Every quarantined trace was + therefore re-attempted on every pass (measured 2026-10-24: 203 files + / 201 unique sids, 200 of them still in the raw queue, 25 already + re-promoted after being quarantined — which is exactly why the leak + went unnoticed). Walk the tree; legacy top-level files count too. + """ + base = os.path.join(cfg.data_dir, "quarantine") + out: set[str] = set() + if not os.path.isdir(base): + return out + for _root, _dirs, files in os.walk(base): + for f in files: + if f.endswith(".json"): + out.add(f[:-5]) + return out + + +def cmd_generate(args) -> int: + """Generate ideal-world references, resuming from the latest on-disk state. + + One idempotent pass: pull new candidates -> improve every raw + trace that has no promoted file yet -> export. Because every stage + keys on disk state (raw cache = queue, improved/ = done set, + per-workdir gate caches = paid judge rounds), a killed run resumes + exactly where it stopped on the next invocation; rerunning after + completion is a no-op that just re-exports. + + Pending set = raw cache (incl. split segments #pN) minus + promoted minus quarantined. Discovery + pull are idempotent + (already-pulled raws skip), so every pass also picks up NEW + sessions; --no-pull makes a pass purely resume the existing queue. + A partial improve pass (some traces fail their gates) still + exports the grown corpus and exits nonzero — the next pass + retries the failures via --skip-existing. + + Quarantine policy: a quarantined sid is SKIPPED by default. The + quarantine classes (echo / rework-echo / rework-language / + rework-reframe) mark traces whose defect the repair layer could not + fix; re-queueing them every pass is pure waste (they re-fail, or + worse, slip back through a later gate change and get re-promoted as + a duplicate). --requeue-quarantined is the explicit opt-in for a + one-time re-queue after a repair-layer change. + """ + cfg = _cfg(args) + if getattr(args, "no_triage", False): + cfg.triage = False + r = MultiSourceReader(cfg) + max_attempts = args.max_attempts + requeue_quarantined = bool(getattr(args, "requeue_quarantined", False)) + quarantined = set() if requeue_quarantined else _quarantined_sids(cfg) + + def _pending_ids() -> list[str]: + if not os.path.isdir(cfg.raw_dir): + return [] + improved = ( + {f[:-5] for f in os.listdir(cfg.improved_dir) + if f.endswith(".json")} + if os.path.isdir(cfg.improved_dir) else set() + ) + return [ + sid for sid in ( + f[:-5] for f in os.listdir(cfg.raw_dir) + if f.endswith(".json") + ) + if sid not in improved and sid not in quarantined + ] + + rc_improve = 0 + if not getattr(args, "no_pull", False): + # discover + pull ONLY when the local queue is empty: mid-batch + # kills resume the backlog first, discovery re-runs when dry. + pending_before = _pending_ids() + if pending_before: + print(f"-- {len(pending_before)} pending raw traces " + f"(resuming backlog)") + sessions = r.list_sessions( + sources=_parse_sources(args), + min_messages=_pick(args, cfg, "min_messages"), + max_messages=_pick(args, cfg, "max_messages"), + min_tool_calls=_pick(args, cfg, "min_tool_calls"), + limit=args.limit, + ) + if sessions: + rc_pull = _pull_sessions(cfg, r, sessions) + if rc_pull != 0: + return rc_pull + elif not pending_before: + print("no candidate sessions found in any store") + + ids = _pending_ids() + if quarantined: + raw_sids = {f[:-5] for f in os.listdir(cfg.raw_dir) + if f.endswith(".json")} if os.path.isdir(cfg.raw_dir) else set() + in_queue = len(quarantined & raw_sids) + print(f"-- {len(quarantined)} quarantined sids skipped " + f"({in_queue} of them still in the raw queue; " + f"--requeue-quarantined to retry)") + if not ids: + print("-- no pending traces; exporting current corpus") + return _generate_export(args) + + # self-reference guard (same layers as cmd_improve; the raw cache + # from a pre-guard run may still hold improver children) + child_ids = r.improver_child_ids() + if child_ids: + drop = [i for i in ids if i in child_ids] + if drop: + ids = [i for i in ids if i not in child_ids] + print(f"-- skipped {len(drop)} improver session(s) " + f"(self-referential data)") + ids = _filter_improver_children_live(cfg, r, ids) + if not ids: + print("-- no improvable traces left after self-reference guard") + return _generate_export(args) + + rc_improve = _improve_ids( + cfg, r, ids, + max_attempts=max_attempts, + workers=args.workers, + skip_existing=True, + ) + if rc_improve != 0: + # partial pass: export what was promoted anyway, keep the + # nonzero rc so the wrapper's loop counts remaining failures + _generate_export(args) + return rc_improve + return _generate_export(args) + + +def _generate_export(args) -> int: + """Export stage shared by cmd_generate (always runs at the end). + + The improve exit code is deliberately not propagated from a + partial pass: --skip-existing makes the NEXT pass retry failures, + and the dataset must still refresh after a partial batch. Errors + (bad exit codes) still surface here. + + Newly promoted traces are labeled before export so the per-topic + split stays complete without a manual `classify` step. + """ + cfg = _cfg(args) + try: + src_dir = cfg.improved_dir + files = [os.path.join(src_dir, f) for f in sorted(os.listdir(src_dir)) + if f.endswith(".json")] + new = _classify_files(cfg, files, workers=8, quiet=True) + if new: + print(f" classified {len(new)} new traces") + except Exception as e: # labeling must never block the export + print(f" classify skipped: {e}") + return cmd_export(types.SimpleNamespace( + config=args.config, out=getattr(args, "out", None), + keep_reasoning=False, + val_frac=getattr(args, "val_frac", 0.0) or 0.0, + force_audit_fail=False, + by_topic=True, + )) + + +def cmd_run(args) -> int: + """One command: find candidates everywhere, pull, improve, export. + + Chain: list (all sources) -> pull -> improve (--skip-existing by default + so reruns are resumable) -> export. The export only runs when the + improve batch is clean (honest exit codes propagate). + """ + cfg = _cfg(args) + r = MultiSourceReader(cfg) + sessions = r.list_sessions( + sources=_parse_sources(args), + min_messages=_pick(args, cfg, "min_messages"), + max_messages=_pick(args, cfg, "max_messages"), + min_tool_calls=_pick(args, cfg, "min_tool_calls"), + limit=args.limit, + ) + if not sessions: + print("no candidate traces found — machine has no improvable sessions") + return 2 + _print_sessions(sessions) + rc_pull = _pull_sessions(cfg, r, sessions) + if rc_pull != 0: + return rc_pull + rc_improve = _improve_ids( + cfg, r, + [s["session_id"] for s in sessions], + max_attempts=args.max_attempts, + workers=args.workers, + skip_existing=True, # run is resumable by default + ) + if rc_improve != 0: + return rc_improve + return cmd_export(types.SimpleNamespace( + config=args.config, out=None, keep_reasoning=False + )) + + +def cmd_timesensitivity(args) -> int: + """Census: does the ideal answer decay with wall-clock time? + + Two axes decide the class: the task asks about an EXTERNAL changing fact, + and the agent had (or lacked) a retrieval tool. A temporal fact answered + from weights is a bad SFT target — the model learns to assert it. The + same fact fetched with `web_search` is a good target: the model learns + the procedure, which does not decay. + """ + from traceforge.timesensitivity import CLASSES, CLASS_LABELS, detect + + cfg = _cfg(args) + records: list = [] + if args.dataset: + path = args.dataset + if not os.path.exists(path): + print(f"error: not found: {path}") + return 2 + with open(path, "r", encoding="utf-8", errors="replace") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + records.append(json.loads(line)) + except json.JSONDecodeError: + continue + if args.limit and len(records) >= args.limit: + break + label = path + else: + if not os.path.isdir(cfg.improved_dir): + print(f"error: no improved traces at {cfg.improved_dir}") + return 2 + files = sorted(f for f in os.listdir(cfg.improved_dir) + if f.endswith(".json")) + if args.limit: + files = files[:args.limit] + for f in files: + try: + records.append(Trace.load(os.path.join(cfg.improved_dir, f))) + except Exception: + continue + label = cfg.improved_dir + if not records: + print("no traces to census") + return 2 + + counts = {c: 0 for c in CLASSES} + rewrite = 0 + examples: list[tuple[str, list[str]]] = [] + for r in records: + d = detect(r) + counts[d["class"]] = counts.get(d["class"], 0) + 1 + if d.get("needs_rewrite"): + rewrite += 1 + if d["class"] == "time_sensitive" and len(examples) < 10: + sid = (getattr(r, "session_id", None) + or ((r.get("meta") or {}).get("traceforge_session_id") + if isinstance(r, dict) else None) or "?") + examples.append((sid, d.get("temporal_facts") or [])) + n = len(records) + census = { + "total": n, + "classes": counts, + "labels": CLASS_LABELS, + "needs_rewrite": rewrite, + "time_sensitive_share": round(counts.get("time_sensitive", 0) / n, 4), + "source": label, + } + if args.cache: + from traceforge.space.stats import _write_json + st = os.stat(args.dataset) if args.dataset else None + _write_json(args.cache, { + "key": ({"mtime": int(st.st_mtime), "size": st.st_size} + if st else {"source": label}), + "census": {**census, "available": True}, + }) + print(f"-- census cache -> {args.cache}") + if args.json: + print(json.dumps(census, indent=2)) + return 0 + print(f"-- time sensitivity over {n} traces ({label})") + for c in CLASSES: + print(f" {c:20s} {counts.get(c, 0):6d} " + f"{100.0 * counts.get(c, 0) / n:5.1f}% {CLASS_LABELS[c]}") + print(f" needs rewrite (retrieval tool available, never called): {rewrite}") + if examples: + print(" examples:") + for sid, facts in examples: + print(f" {sid[:48]:48s} {', '.join(facts)}") + return 0 + + +def cmd_space(args) -> int: + """Portable web app: submit traces, sponsor datasets, public dashboard.""" + from traceforge.space import (SpacePipeline, TraceStore, fingerprint, + parse_records, read_upload, space_config) + + if args.action == "serve": + from traceforge.space.app import main as app_main + argv: list[str] = [] + if args.host: + argv += ["--host", args.host] + if args.port: + argv += ["--port", str(args.port)] + if args.share: + argv += ["--share"] + return app_main(argv) + + cfg = space_config(data_dir=args.data_dir) + store = TraceStore(cfg.data_dir) + + if args.action == "hash": + if not args.files: + print("error: pass one or more trace files") + return 2 + for path in args.files: + records, errors = read_upload(path) + for e in errors: + print(f" {path}: {e}") + for rec in records: + fp = fingerprint(rec) + print(f"{fp['hash']} {os.path.basename(path)} " + f"{fp['n_messages']} msgs · {fp['n_tool_calls']} calls · " + f"{fp['language']} · {fp['task'][:60]}") + return 0 + + if args.action == "stats": + d = store.stats() + if args.json: + print(json.dumps(d, indent=2)) + return 0 + print(f"-- space store: {store.root}") + print(f" seeds {d['seeds']} · done {d['done']} · failed {d['failed']} " + f"· queued {d['queued']} · running {d['running']}") + for s in d["sponsors"]: + print(f" sponsor {s['sponsor']:24s} {s['seeds']:5d} seeds " + f"({s['done']} done, {s['failed']} failed)") + if d["time_sensitivity"]: + print(" time sensitivity: " + ", ".join( + f"{k} {v}" for k, v in d["time_sensitivity"].items())) + return 0 + + if args.action == "process": + if not args.files: + print("error: pass one or more trace files") + return 2 + pipeline = SpacePipeline(store, cfg) + rc = 0 + for path in args.files: + records, errors = read_upload(path) + for e in errors: + print(f" {path}: {e}") + rc = 1 + for rec in records: + res = pipeline.submit(rec, sponsor=args.sponsor) + h = res["hash"] + if res.get("duplicate") and res.get("status") == "done": + print(f"DUPLICATE {h} — existing output returned") + continue + out = pipeline.process(h) + print(f"{out.get('status', '?').upper():9s} {h} " + f"{out.get('duration_s', 0):.0f}s " + f"{out.get('error', '')[:200]}") + if out.get("status") != "done": + rc = 1 + return rc + + print(f"error: unknown action {args.action}") + return 2 + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + prog="traceforge", + description="Evaluation harness for agent models: score a model by how " + "little an ideal-world rewrite has to edit its real traces", + ) + p.add_argument("--config", help="YAML config overrides") + sub = p.add_subparsers(dest="cmd", required=True) + + common = argparse.ArgumentParser(add_help=False) + common.add_argument("--min-messages", type=int, default=None) + common.add_argument("--max-messages", type=int, default=None) + common.add_argument("--min-tool-calls", type=int, default=None) + common.add_argument("--limit", type=int, default=None) + common.add_argument( + "--source", "-s", default="all", + help="trace source filter: all | hermes | claude | comma list " + "(default: all)", + ) + + sp = sub.add_parser( + "stats", help="trace-store counts (all sources)", parents=[common] + ) + sp.set_defaults(fn=cmd_stats) + + sp = sub.add_parser( + "list", help="list candidate traces (all sources)", parents=[common] + ) + sp.set_defaults(fn=cmd_list) + + sp = sub.add_parser( + "pull", help="normalize candidate traces to data/traces/raw/", parents=[common] + ) + sp.set_defaults(fn=cmd_pull) + + sp = sub.add_parser("improve", help="run the agentic improver on traces") + sp.add_argument("session_ids", nargs="*", help="session ids (empty = --from-raw)") + sp.add_argument("--from-raw", action="store_true", help="improve all pulled traces") + sp.add_argument("--max-attempts", type=int, default=2) + sp.add_argument( + "--workers", "-w", type=int, default=1, + help="concurrent improver sessions = parallel requests to the LLM " + "endpoint (real batch size; default 1 = sequential)", + ) + sp.add_argument( + "--skip-existing", action="store_true", + help="skip traces that already have data/traces/improved/.json", + ) + sp.add_argument( + "--no-triage", action="store_true", + help="skip the LLM triage pass even when enabled in config", + ) + sp.set_defaults(fn=cmd_improve) + + sp = sub.add_parser( + "triage", + help="label traces (keep/drop/rephrase/fact_check/summarize) via " + "structured generation; prints the route per trace", + ) + sp.add_argument("session_ids", nargs="*", help="session ids (empty = --from-raw)") + sp.add_argument("--from-raw", action="store_true", help="triage all pulled traces") + sp.add_argument( + "--force", action="store_true", + help="re-label even when data/triage/.json exists", + ) + sp.set_defaults(fn=cmd_triage) + + sp = sub.add_parser( + "funnel", + help="funnel census: where do traces drop before promotion? (read-only)", + ) + sp.add_argument( + "--sample", type=int, default=200, + help="artifact sample size for the gate re-adjudication (default 200)", + ) + sp.set_defaults(fn=cmd_funnel) + + sp = sub.add_parser( + "benchmark", + help="score a model's traces by how little the pipeline had to edit " + "(fidelity + triage labels -> 0-100); read-only", + ) + sp.add_argument("--raw-dir", default=None, + help="model traces (default data/traces/raw/)") + sp.add_argument("--improved-dir", default=None, + help="ideal rewrites (default data/traces/improved/)") + sp.add_argument("--triage-dir", default=None, + help="triage label files (default data/triage/)") + sp.add_argument("--label", default=None, + help="name this run (model/checkpoint); appends a row to " + "the leaderboard next to the report") + sp.add_argument("--out", default=None, + help="report path (default data/benchmark/report.json)") + sp.add_argument("--json", action="store_true", + help="print the JSON report instead of the table") + sp.add_argument("--limit", type=int, default=None, + help="score at most N pairs (smoke test)") + sp.add_argument("--source", "-s", default="all", + help="filter by the model trace's source: all | hermes | " + "claude | codex | comma list") + sp.add_argument("--w-similarity", type=float, default=0.6, + help="weight of fidelity in the composite (default 0.6)") + sp.add_argument("--w-labels", type=float, default=0.4, + help="weight of the triage keep rate (default 0.4)") + sp.set_defaults(fn=cmd_benchmark) + + sp = sub.add_parser( + "regate", + help="adjudicate existing workdir artifacts with the windowed " + "judge (gates locally, no agent runtime); promotes on PASS", + ) + sp.add_argument("session_ids", nargs="*", + help="session ids (empty = all workdir artifacts)") + sp.add_argument("--workers", "-w", type=int, default=8) + sp.set_defaults(fn=cmd_regate) + + sp = sub.add_parser( + "crossexam", + help="re-judge PROMOTED traces with a second judge family; " + "disagreements are reported for human review (read-only)", + ) + sp.add_argument("session_ids", nargs="*", + help="session ids (empty = whole promoted corpus)") + sp.add_argument("--model", "-m", required=True, + help="cross judge model (DIFFERENT family than the " + "improver, e.g. gemma4-31b-selfhosted)") + sp.add_argument("--workers", "-w", type=int, default=8) + sp.set_defaults(fn=cmd_crossexam) + + sp = sub.add_parser( + "audit", + help="mechanical corpus audit (self-talk, clip markers, " + "near-duplicates); zero LLM calls", + ) + sp.add_argument("files", nargs="*") + sp.add_argument("-v", "--verbose", action="store_true") + sp.set_defaults(fn=cmd_audit) + + sp = sub.add_parser("validate", help="structurally validate improved traces") + sp.add_argument("files", nargs="*") + sp.set_defaults(fn=cmd_validate) + + sp = sub.add_parser( + "verify", + help="web fact-check improved traces' external claims (adversarial verifier)", + ) + sp.add_argument("files", nargs="*") + sp.set_defaults(fn=cmd_verify) + + sp = sub.add_parser( + "classify", + help="label improved traces with a task/topic bucket (for per-topic export)", + ) + sp.add_argument("files", nargs="*") + sp.add_argument("--heuristic-only", action="store_true", + help="skip the LLM tier (free, deterministic, coarser)") + sp.add_argument("--force", action="store_true", + help="reclassify traces that already carry a task label") + sp.add_argument("--workers", "-w", type=int, default=8, + help="concurrent LLM calls (default 8)") + sp.set_defaults(fn=cmd_classify) + + sp = sub.add_parser("export", help="export references to training JSONL " + "(downstream reuse, not the repo's scope)") + sp.add_argument("--out", default=None) + sp.add_argument("--keep-reasoning", action="store_true", + help="include per-message reasoning in the export") + sp.add_argument("--force-audit-fail", action="store_true", + help="export even when the mechanical audit fails") + sp.add_argument("--val-frac", type=float, default=0.0, + help="hold out a deterministic fraction of traces as " + "val.jsonl (stable session-id hash split)") + sp.add_argument("--by-topic", action="store_true", + help="also write one train..jsonl per task bucket " + "(labels from `traceforge classify`; unlabeled traces " + "are classified on the fly with the heuristic)") + sp.add_argument("--drop-time-sensitive", action="store_true", + help="drop traces whose answer is a temporal fact the " + "agent had no retrieval tool for (the answer decays; " + "training on it teaches the model to assert it)") + sp.set_defaults(fn=cmd_export) + + sp = sub.add_parser( + "kto", + help="build the preference (KTO) corpus: rollout seeds, then " + "length-matched binary labels", + ) + sp.add_argument("--prompts", action="store_true", + help="extract rollout seeds from the SFT corpus " + "(the default action)") + sp.add_argument("--build", default=None, metavar="ROLLOUTS", + help="labelled rollouts JSONL -> length-matched KTO rows") + sp.add_argument("--provisional", action="store_true", + help="build a NOT-on-policy smoke-test set: positives from " + "the SFT corpus, negatives from triage-dropped turns " + "(replace with rollout-built rows after SFT)") + sp.add_argument("--out", default=None) + sp.add_argument("--max-prompts", type=int, default=0, + help="cap the number of rollout seeds (0 = all)") + sp.add_argument("--turns", choices=("final", "all"), default="final", + help="one seed per trace at the final assistant turn " + "(default, verifiable) or every assistant turn") + sp.add_argument("--max-prefix-chars", type=int, default=0, + help="bound each seed's prefix (0 = full context); keeps " + "the opening task and the most recent turns") + sp.add_argument("--max-len-ratio", type=float, default=1.5, + help="drop a negative longer than this multiple of the " + "shortest desirable completion (default 1.5)") + sp.set_defaults(fn=cmd_kto) + + sp = sub.add_parser( + "run", help="one command: find all traces, pull, improve, export", + parents=[common], + ) + sp.add_argument("--max-attempts", type=int, default=2) + sp.add_argument( + "--workers", "-w", type=int, default=1, + help="concurrent improver sessions = parallel requests to the LLM " + "endpoint (real batch size; default 1 = sequential)", + ) + sp.set_defaults(fn=cmd_run) + + sp = sub.add_parser( + "generate", + help="generate data, resuming from the latest on-disk state; " + "rerun after any interruption to continue", + parents=[common], + ) + sp.add_argument("--max-attempts", type=int, default=2) + sp.add_argument( + "--workers", "-w", type=int, default=1, + help="concurrent improver sessions = parallel requests to the LLM " + "endpoint (real batch size; default 1 = sequential)", + ) + sp.add_argument("--no-pull", action="store_true", + help="do not discover/pull new sessions; only improve " + "raws already in data/traces/raw/") + sp.add_argument("--no-triage", action="store_true", + help="skip the LLM triage pass even when enabled in config") + sp.add_argument( + "--requeue-quarantined", action="store_true", + help="re-queue traces parked in data/quarantine// (default: " + "skip them). Use ONCE after a repair-layer change; a quarantined " + "trace otherwise re-fails on every pass", + ) + sp.add_argument("--out", default=None, + help="export path (default data/dataset/train.jsonl)") + sp.add_argument("--val-frac", type=float, default=0.0, + help="hold out a deterministic fraction of traces as " + "val.jsonl (stable session-id hash split)") + sp.set_defaults(fn=cmd_generate) + + sp = sub.add_parser( + "timesensitivity", + help="census: which ideal answers decay with wall-clock time " + "(temporal fact, no retrieval tool)", + ) + sp.add_argument("--dataset", default=None, + help="JSONL export to census (default: data/traces/improved/)") + sp.add_argument("--limit", type=int, default=0, + help="cap the number of traces (0 = all)") + sp.add_argument("--json", action="store_true", + help="machine-readable census") + sp.add_argument("--cache", default=None, + help="write the dashboard census cache to this path") + sp.set_defaults(fn=cmd_timesensitivity) + + sp = sub.add_parser( + "space", + help="portable web app: submit traces, sponsor datasets, public dashboard", + ) + sp.add_argument("action", choices=("serve", "hash", "stats", "process")) + sp.add_argument("files", nargs="*", help="trace files (hash / process)") + sp.add_argument("--host", default=None) + sp.add_argument("--port", type=int, default=None) + sp.add_argument("--share", action="store_true") + sp.add_argument("--data-dir", default=None, + help="store root (default data/space)") + sp.add_argument("--sponsor", default="", help="sponsor name for `process`") + sp.add_argument("--json", action="store_true") + sp.set_defaults(fn=cmd_space) + + return p + + +def main(argv=None) -> int: + args = build_parser().parse_args(argv) + try: + return args.fn(args) + except FileNotFoundError as e: + print(f"error: {e}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file