diff --git "a/traceforge/cli.py" "b/traceforge/cli.py" --- "a/traceforge/cli.py" +++ "b/traceforge/cli.py" @@ -1,2766 +1,3039 @@ -"""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: - # The SAME pipeline the Gradio app runs. One implementation, so - # a gate, a no-op rule or a measurement can never drift between - # the batch and the Space again. - from traceforge.improver import ImproverResult - from traceforge.space.pipeline import process_trace - - pres = process_trace(cfg, t, max_attempts=max_attempts, - seed=t, label_task=False, - reset_workdir=False, status_cb=cb) - res = ImproverResult( - session_id=sid, ok=pres.ok, improved=pres.improved, - reply=pres.route, errors=pres.errors, - attempts=pres.attempts, duration_s=pres.duration_s, - ) - if res.ok and res.improved is not None: - out = os.path.join(cfg.improved_dir, f"{sid}.json") - res.improved.save(out) - note = (" · no change (the raw trace already reads as ideal)" - if pres.no_change else "") - if pres.route == "regate": - prog.write(f"IMPROVED {sid} (regate, " - f"{res.duration_s:.0f}s){note}") - else: - prog.write( - f"IMPROVED {sid} ({res.attempts} attempts, " - f"{res.duration_s:.0f}s, {res.reply}){note}" - ) - else: - prog.write(f"FAILED {sid}: {res.errors}") - row = res.to_dict() - row["route"] = pres.route - row["no_change"] = pres.no_change - row["time_sensitivity"] = pres.time_sensitivity - _write_manifest_row(cfg, row) - if liv is not None: - liv.registry.end(sid) - return sid, res - # Legacy agent engine (`splice: false`). It is not the pipeline the - # Space runs, so it keeps its own behaviour; the batch only reaches - # it when the config explicitly turns the splice engine off. - 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__": +"""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). + + The traceback is recorded, not just the message. A crash row that says + only "[Errno 22] Invalid argument" is undebuggable after the fact: the + batch keeps running, the log scrolls, and the one piece of evidence + that would locate the fault is gone. Measured 2026-10-28: 29 such rows + and no way to tell which call raised. + """ + import traceback as _tb + + return { + "session_id": sid, "ok": False, + "errors": [f"worker crashed: {exc}"], + "crash_v": 2, + "traceback": _tb.format_exc()[-4000:], + "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: + # The SAME pipeline the Gradio app runs. One implementation, so + # a gate, a no-op rule or a measurement can never drift between + # the batch and the Space again. + from traceforge.improver import ImproverResult + from traceforge.space.pipeline import process_trace + + pres = process_trace(cfg, t, max_attempts=max_attempts, + seed=t, label_task=False, + reset_workdir=False, status_cb=cb) + res = ImproverResult( + session_id=sid, ok=pres.ok, improved=pres.improved, + reply=pres.route, errors=pres.errors, + attempts=pres.attempts, duration_s=pres.duration_s, + ) + if res.ok and res.improved is not None: + out = os.path.join(cfg.improved_dir, f"{sid}.json") + res.improved.save(out) + note = (" · no change (the raw trace already reads as ideal)" + if pres.no_change else "") + if pres.route == "regate": + prog.write(f"IMPROVED {sid} (regate, " + f"{res.duration_s:.0f}s){note}") + else: + prog.write( + f"IMPROVED {sid} ({res.attempts} attempts, " + f"{res.duration_s:.0f}s, {res.reply}){note}" + ) + else: + prog.write(f"FAILED {sid}: {res.errors}") + row = res.to_dict() + row["route"] = pres.route + row["no_change"] = pres.no_change + row["time_sensitivity"] = pres.time_sensitivity + _write_manifest_row(cfg, row) + if liv is not None: + liv.registry.end(sid) + return sid, res + # Legacy agent engine (`splice: false`). It is not the pipeline the + # Space runs, so it keeps its own behaviour; the batch only reaches + # it when the config explicitly turns the splice engine off. + 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 _load_harness_map(cfg) -> dict[str, str]: + """sid -> harness family, cached by scripts/harness_census.py. + + Reading it from the raws costs one file open per trace (41k opens, ~10 + min on Windows), which is far too slow to do inside the queue sort. The + census writes the map once; a missing map just means the harness + tiebreaker is inert, not that the pass fails. + """ + path = os.path.join(os.path.dirname(cfg.manifest_path), ".harness_map.json") + try: + with open(path, encoding="utf-8") as fh: + data = json.load(fh) + except (OSError, ValueError): + return {} + return {str(k): str(v) for k, v in data.items()} if isinstance(data, dict) else {} + + +def _harness_priority(cfg) -> dict[str, int]: + """Rank harness families by how far BELOW their target share they are. + + Reads the balance report (refreshed by every export) so the queue follows + the policy in config/balance.yaml rather than a hardcoded order. Falls + back to the deployment order when no report exists yet. + """ + fallback = {"hermes": 0, "codex": 1, "claude-code": 2, "nemotron": 3} + path = os.path.join(os.path.dirname(cfg.manifest_path), "balance", "report.json") + try: + with open(path, encoding="utf-8") as fh: + rep = json.load(fh) + rows = rep["axes"]["harness"]["rows"] + except (OSError, ValueError, KeyError, TypeError): + return fallback + # Most under-target first; `other` is unknown provenance, never a target. + rows = [r for r in rows if r.get("key") != "other"] + rows.sort(key=lambda r: r.get("delta_n", 0)) + return {r["key"]: i for i, r in enumerate(rows)} + + +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] + + # Harness deficit: the queue ranks by class and size, which is right for + # throughput but blind to the mix. Measured 2026-10-29: hermes is 17.4% + # of the corpus against a 32% target while nemotron is 19.5% against 2%, + # so promoting in size order deepens the imbalance. Ranked AFTER class + # (a low-yield class must not jump the queue) and BEFORE size (so the + # under-target family is worked through in cheap-first order). + _hprio = _harness_priority(cfg) + _hmap = _load_harness_map(cfg) + + def _harness_rank(sid: str) -> int: + hs = _hmap.get(sid) + if hs is None: + return len(_hprio) # unknown provenance: neutral tail + return _hprio.get(hs, len(_hprio)) + + ids = sorted(ids, key=lambda s: (_attempted(s), _class_rank(s), + _harness_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, + keep_reasoning: bool = False) -> 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 kept per message (keep_reasoning=True)" + if keep_reasoning else + "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}") + # Balance report: the dashboard reads this file (it must not import + # traceforge), so every export refreshes it. Measuring the traces that + # were just written keeps the report and the shipped mix in lockstep. + try: + from traceforge import balance as bal + policy = bal.BalancePolicy.load(bal.POLICY_PATH) + labels = bal.load_labels(os.path.join(os.path.dirname(out), + "..", "tasks.jsonl")) + cen = bal.census(traces, labels) + rep = bal.report(cen, policy) + bal.write_report(rep, os.path.join(os.path.dirname(out), "..", + "balance", "report.json"), + extra={"plan": bal.plan(rep, policy), + "fill": bal.fill_to(cen, policy), + "grow": bal.converge(cen, policy)}) + print(f" balance: score {rep['balance_score']:.3f} " + f"(1.0 = on target) -> data/balance/report.json") + except Exception as e: # bookkeeping must never block an export + print(f" balance: report skipped ({e})") + # 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_balance(args) -> int: + """Report the corpus mix against the target policy, and plan the gap. + + Reads the exported dataset (the shipped artifact) rather than improved/, + because the export is what the trainer sees: dedup, time-sensitivity and + any balanced selection have already been applied there. + """ + from traceforge import balance as bal + + cfg = _cfg(args) + policy = bal.BalancePolicy.load( + getattr(args, "policy", None) or bal.POLICY_PATH) + path = getattr(args, "dataset", None) or os.path.join( + cfg.data_dir, "dataset", "train.jsonl") + if not os.path.exists(path): + print(f"no dataset at {path} — run `traceforge export` first") + return 2 + labels = bal.load_labels(os.path.join(cfg.data_dir, "tasks.jsonl")) + limit = getattr(args, "limit", 0) or 0 + records = [] + with open(path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + records.append(json.loads(line)) + if limit and len(records) >= limit: + break + c = bal.census(records, labels) + rep = bal.report(c, policy) + fill = bal.fill_to(c, policy, getattr(args, "target_n", 0) or None) + grow = bal.converge(c, policy) + if getattr(args, "json", False): + payload = {"report": rep, "fill": fill, "grow": grow} + if getattr(args, "plan", False): + payload["plan"] = bal.plan(rep, policy) + print(json.dumps(payload, indent=1, ensure_ascii=False)) + else: + print(bal.format_report(rep, top=getattr(args, "top", 0) or 0)) + if getattr(args, "plan", False): + print() + print(bal.format_plan(bal.plan(rep, policy), + limit=getattr(args, "top", 0) or 20, + fill=fill, grow=grow)) + if getattr(args, "write", False): + out = getattr(args, "out", None) or os.path.join( + cfg.data_dir, "balance", "report.json") + bal.write_report(rep, out, extra={ + "plan": bal.plan(rep, policy), + "fill": fill, + "grow": grow, + }) + if not getattr(args, "json", False): + print(f"\n wrote {out}") + return 0 + + +def cmd_synth(args) -> int: + """Synthetic task generation: plan from the balance gap, write tasks, + execute them with a real agent, ingest the traces into the pool.""" + from traceforge import synth + + cfg = _cfg(args) + action = getattr(args, "action", None) or "status" + if action == "plan": + return synth.cli_plan(cfg, args) + if action == "tasks": + return synth.cli_tasks(cfg, args) + if action == "run": + return synth.cli_run(cfg, args) + if action == "ingest": + return synth.cli_ingest(cfg, args) + return synth.cli_status(cfg, args) + + +def cmd_export(args) -> int: + cfg = _cfg(args) + src_dir = cfg.improved_dir + # Reasoning channel: the rewriter WRITES reasoning (the source harnesses + # only record a summary), and the distill stage puts a trace's dead ends + # there. `generate`'s export already honours cfg.export_reasoning; this + # command defaulted to False, so a manual/cron `traceforge export + # --by-topic` silently shipped a corpus with no reasoning at all + # (measured 2026-10-29: 0 of 3000 rows in data/dataset/train.jsonl). + # Resolve the flag against the config instead of hard-defaulting off. + keep_reasoning = getattr(args, "keep_reasoning", None) + if keep_reasoning is None: + keep_reasoning = bool(getattr(cfg, "export_reasoning", True)) + # 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") + # Balance: the cheap half of rebalancing is capping the over-represented + # cells, so the SHIPPED mix can be correct today while generation fills + # the under-represented ones over time. The full corpus stays on disk. + if getattr(args, "balance", False): + from traceforge import balance as bal + policy = bal.BalancePolicy.load( + getattr(args, "balance_policy", None) or bal.POLICY_PATH) + labels = bal.load_labels(os.path.join(cfg.data_dir, "tasks.jsonl")) + raw_axes = getattr(args, "balance_axes", None) or "language,task,harness" + axes = tuple(a.strip() for a in raw_axes.split(",") if a.strip()) + traces, bstats = bal.select_balanced(traces, policy, labels, axes=axes) + print(f"-- balance: kept {bstats['out']}/{bstats['in']} trace(s), " + f"dropped {bstats['dropped_total']} from over-target cells " + f"(axes {','.join(axes)}; policy {policy.source})") + for cell, k in list(bstats["dropped"].items())[:8]: + print(f" {cell:<22s} -{k}") + 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=keep_reasoning) + n = export_chat_jsonl( + train_traces, out, + keep_reasoning=keep_reasoning) + 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, keep_reasoning) + 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), + keep_reasoning=keep_reasoning) + return 0 + n = export_chat_jsonl( + traces, out, keep_reasoning=keep_reasoning + ) + print(f"exported {n} traces -> {out}") + if getattr(args, "by_topic", False): + counts = _export_by_topic( + traces, out, cfg, keep_reasoning) + 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), + keep_reasoning=keep_reasoning) + 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=getattr(cfg, "export_reasoning", True), + 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=None + )) + + +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", dest="keep_reasoning", + action="store_true", default=None, + help="include per-message reasoning in the export " + "(default: the config's export_reasoning)") + sp.add_argument("--no-reasoning", dest="keep_reasoning", + action="store_false", + help="drop per-message reasoning even when the config " + "keeps it") + 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.add_argument("--balance", action="store_true", + help="deterministically downsample over-target " + "(language, task) cells to the config/balance.yaml " + "target mix; the full corpus stays on disk") + sp.add_argument("--balance-policy", metavar="PATH", + help="balance policy YAML (default config/balance.yaml)") + sp.add_argument( + "--balance-axes", metavar="LIST", default="language,task,harness", + help="which caps --balance applies (default: all three). " + "'harness' alone fixes the harness mix without the far more " + "aggressive language/task caps.", + ) + sp.set_defaults(fn=cmd_export) + + sp = sub.add_parser( + "balance", + help="report the corpus mix against the target policy and plan the gap", + ) + sp.add_argument("--dataset", default=None, + help="exported JSONL to measure (default " + "data/dataset/train.jsonl)") + sp.add_argument("--config", default=None, + help="traceforge config YAML (pins provider/model; " + "default: built-in defaults)") + sp.add_argument("--policy", default=None, metavar="PATH", + help="balance policy YAML (default config/balance.yaml)") + sp.add_argument("--plan", action="store_true", + help="also print the generation work order (cells under " + "target, largest deficit first)") + sp.add_argument("--top", type=int, default=0, + help="show only the first N rows per axis (0 = all)") + sp.add_argument("--target-n", type=int, default=0, + help="plan the cap+fill for this final corpus size " + "(0 = the current size, i.e. balance what we have)") + sp.add_argument("--limit", type=int, default=0, + help="measure only the first N records (0 = all)") + sp.add_argument("--json", action="store_true") + sp.add_argument("--write", action="store_true", + help="write the report JSON for the dashboard") + sp.add_argument("--out", default=None, + help="report path (default data/balance/report.json)") + sp.set_defaults(fn=cmd_balance) + + sp = sub.add_parser( + "synth", + help="generate synthetic traces from scratch to fill balance gaps: " + "plan -> tasks -> run -> ingest", + ) + sp.add_argument("action", nargs="?", default="status", + choices=("status", "plan", "tasks", "run", "ingest"), + help="status (default) | plan | tasks | run | ingest") + sp.add_argument("--policy", default=None, metavar="PATH", + help="balance policy YAML (default config/balance.yaml)") + sp.add_argument("--config", default=None, + help="traceforge config YAML (pins provider/model; " + "default: built-in defaults)") + sp.add_argument("--dataset", default=None, + help="exported JSONL the gap is measured against") + sp.add_argument("--cells", type=int, default=0, + help="plan/tasks: only the first N cells of the work order") + sp.add_argument("--per-cell", type=int, default=1, + help="tasks: how many task descriptions per cell") + sp.add_argument("--limit", type=int, default=0, + help="run: cap the number of tasks executed this pass") + sp.add_argument("--retry-failed", action="store_true", + help="run: also retry tasks that failed previously") + sp.add_argument("--workers", "-w", type=int, default=1, + help="run: concurrent agent sessions") + sp.add_argument("--max-turns", type=int, default=40, + help="run: agent turn cap per task (bounds cost)") + sp.add_argument("--timeout", type=int, default=1800, + help="run: wall-clock cap per task, seconds") + sp.add_argument("--model", default=None, + help="run: model id (default: the config's model)") + sp.add_argument("--provider", default=None, + help="run: provider block (default: the config's provider)") + sp.add_argument("--json", action="store_true") + sp.set_defaults(fn=cmd_synth) + + 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