Something went wrong. Try again.
Parametric learning in your code harness.
Something went wrong. Try again.
20 kB · 445 lines
TypeScript
at letta-code
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "@earendil-works/pi-coding-agent";import { readFileSync } from "node:fs";import path from "node:path";import { ensureMachineDataDir, machineDataFile, readMachineSettings, writeMachineSettings,} from "../lib/storage.ts";import { countJsonlRows, flagOn, parseOptions, rel, shellSplit } from "../lib/cli.ts";import { appendTaggedExample, formatTaggedExampleList } from "../lib/examples.ts";import { forkObservedRun, formatMachineDataStatus } from "../lib/fork.ts";import { formatMachineStatus, joinMachineStatusLine, resolveLastTurnStatus, type MachineStatusLastSave,} from "../lib/statusline.ts";import { CHECKPOINT_STATE_PATH, QWEN_38_27B_256K_MODEL, isTinkerNativeId } from "../lib/tinker-catalog.ts";import { registerMachineReporter, setMachineReportMirror } from "../lib/report.ts";import { applyAssistantRewrite, extractMessageText, lastTaggedTurn, lastUserTurnEntries, previewText, recordsFromSessionEntries, summarizeTrajectory, trajectoryToMessages, type SessionMessageLike, type TaggedExampleTag,} from "../lib/trajectory.ts";
// The 256K peft Inkling catalog entry; /train in train.ts defaults to the// 64K base (DEFAULT_MODEL). They differ on purpose: evals and training run// cheaper at 64K, while interactive student sessions want the long context.const DEFAULT_FT_MODEL = "thinkingmachines/Inkling-Small:peft:262144";const CHECKPOINT_STATE = CHECKPOINT_STATE_PATH;
function resolveTaggedExamplePath(cwd: string, tag: TaggedExampleTag, out?: string, trace = false): string { if (out) return path.resolve(cwd, out); ensureMachineDataDir(cwd); if (tag === "observed") return machineDataFile(cwd, "observed"); if (tag === "bad") return machineDataFile(cwd, "rejected"); return machineDataFile(cwd, trace ? "traces" : "train");}
type ModelRef = { provider: string; id: string };
function describeModel(model?: { provider?: string; id?: string } | null): string { if (!model?.provider || !model.id) return "(none)"; return `${model.provider}/${model.id}`;}
function parseModelRef(raw: string): ModelRef | undefined { const value = raw.trim(); if (!value) return undefined; if (isTinkerNativeId(value)) return { provider: "tinker", id: value }; const slash = value.indexOf("/"); if (slash <= 0) return { provider: "tinker", id: value }; return { provider: value.slice(0, slash), id: value.slice(slash + 1) };}
function latestCheckpointId(): string | undefined { try { const parsed = JSON.parse(readFileSync(CHECKPOINT_STATE, "utf8")) as { checkpoints?: Array<{ id?: string }> }; return parsed.checkpoints?.slice(-1)[0]?.id; } catch { return undefined; }}
export default function (pi: ExtensionAPI) { const sendReport = registerMachineReporter(pi); let workModel: ModelRef | undefined; let lastSave: MachineStatusLastSave | undefined;
function refreshStatus(ctx: ExtensionContext | ExtensionCommandContext) { if (!ctx.hasUI) return; ctx.ui.setStatus("machine-auto", undefined); const entries = ctx.sessionManager.buildContextEntries() as SessionMessageLike[]; const turn = lastTaggedTurn(entries); const lastTurn = resolveLastTurnStatus({ hasTurn: Boolean(turn && !turn.user.trim().startsWith("/")), label: turn?.assistantId ? ctx.sessionManager.getLabel(turn.assistantId) : undefined, lastSave, assistantId: turn?.assistantId, }); ctx.ui.setStatus("machine", joinMachineStatusLine( formatMachineStatus({ autoTrace: readMachineSettings(ctx.cwd).autoTrace, student: ctx.model?.provider === "tinker", lastTurn, train: countJsonlRows(machineDataFile(ctx.cwd, "train")), observed: countJsonlRows(machineDataFile(ctx.cwd, "observed")), }), (tone, text) => ctx.ui.theme.fg(tone, text), )); }
pi.on("session_start", (_event, ctx) => { setMachineReportMirror(!ctx.hasUI); refreshStatus(ctx); });
pi.on("turn_end", (_event, ctx) => { refreshStatus(ctx); });
pi.on("model_select", (_event, ctx) => { refreshStatus(ctx); });
async function saveTaggedTurn(tag: TaggedExampleTag, rawArgs: string, ctx: ExtensionCommandContext) { const { positional, options } = parseOptions(shellSplit(rawArgs)); const traceFlag = flagOn(options, "trace"); const trajectoryFlag = flagOn(options, "trajectory"); const reasoningFlag = flagOn(options, "reasoning"); const note = String(options.note ?? [traceFlag.leftover, trajectoryFlag.leftover, reasoningFlag.leftover, ...positional].filter(Boolean).join(" ")).trim(); const replacement = String(options.as ?? options.rewrite ?? "").trim(); const trace = traceFlag.on || trajectoryFlag.on; const includeReasoning = reasoningFlag.on; const entries = ctx.sessionManager.buildContextEntries() as SessionMessageLike[]; const turn = lastTaggedTurn(entries); if (!turn) { sendReport(`Save ${tag} example`, "No user/assistant turn found to tag. Chat first, then run `/good` or `/bad`.", "warning"); return; } const records = applyAssistantRewrite( recordsFromSessionEntries(lastUserTurnEntries(entries)).filter((record) => includeReasoning || record.role !== "reasoning"), replacement, ); const assistant = replacement || turn.assistant; const messages = trace ? trajectoryToMessages(records) : [ { role: "user", content: turn.user }, { role: "assistant", content: assistant }, ]; const outFile = resolveTaggedExamplePath(ctx.cwd, tag, typeof options.out === "string" ? options.out : undefined, trace); try { const saved = await appendTaggedExample({ tag, messages, records: trace ? records : undefined, note: note || undefined, outFile, sessionFile: ctx.sessionManager.getSessionFile(), userId: turn.userId, assistantId: turn.assistantId, kind: trace ? "trace" : "final", }); if (turn.assistantId) pi.setLabel(turn.assistantId, tag); if (turn.userId && !ctx.sessionManager.getLabel(turn.userId)) pi.setLabel(turn.userId, tag); if (tag === "good" || tag === "bad") lastSave = { assistantId: turn.assistantId, tag, trace }; refreshStatus(ctx); sendReport(`Saved ${tag} ${trace ? "trace" : "example"}`, [ `Appended the last turn to \`${rel(ctx.cwd, outFile)}\` (${saved.count} row${saved.count === 1 ? "" : "s"}).`, saved.row.id ? `id: \`${String(saved.row.id)}\`` : "", trace ? "Reconstructed the full user→tool-call→result→answer trajectory." : "Saved the final assistant answer only.", replacement ? "Used the rewritten assistant text you provided." : "", note ? `Note: ${note}` : "", "", trace ? `## Trajectory\n${summarizeTrajectory(records)}` : `## User\n${previewText(turn.user, 600)}\n\n## Assistant\n${previewText(assistant, 600)}`, "", tag === "good" ? (trace ? "`messages` is Tinker chat JSONL including tool calls. Train with `TrainOnWhat.ALL_ASSISTANT_MESSAGES` so results stay in context but are not the target. `trajectory` keeps the Letta-style records." : "This file is Tinker chat JSONL. Inspect it, then run `/train validate`.") : "Rejected rows stay out of SFT. Pair them later for DPO, or rewrite and `/good --as \"...\"`.", ].filter(Boolean).join("\n"), "success"); } catch (error) { sendReport(`Save ${tag} example failed`, error instanceof Error ? error.message : String(error), "error"); } }
async function listTaggedExamples(rawArgs: string, ctx: ExtensionCommandContext) { const { options } = parseOptions(shellSplit(rawArgs)); const files = [ typeof options.out === "string" ? path.resolve(ctx.cwd, options.out) : undefined, machineDataFile(ctx.cwd, "train"), machineDataFile(ctx.cwd, "traces"), machineDataFile(ctx.cwd, "observed"), machineDataFile(ctx.cwd, "rejected"), ].filter((file, idx, all): file is string => Boolean(file) && all.indexOf(file) === idx); sendReport("Tagged examples", formatTaggedExampleList(ctx.cwd, files)); }
pi.registerCommand("good", { description: "Save the last turn as a good training example", handler: async (args, ctx) => { await saveTaggedTurn("good", args, ctx); }, });
pi.registerCommand("bad", { description: "Save the last turn as a rejected / bad example", handler: async (args, ctx) => { await saveTaggedTurn("bad", args, ctx); }, });
pi.registerCommand("tagged", { description: "List recent judged examples and automatic observations", handler: async (args, ctx) => { await listTaggedExamples(args, ctx); }, });
async function switchLane(lane: "work" | "student", rawArgs: string, ctx: ExtensionCommandContext) { const current = ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : undefined; if (lane === "work") { if (current && current.provider === "tinker") workModel = workModel ?? current; const target = parseModelRef(rawArgs) ?? workModel; if (!target) { sendReport("Work model", "No daily-driver model remembered yet. Stay on Grok (or whatever you use), then `/machine student`. The next `/machine work` will jump back.", "warning"); return; } if (current && current.provider !== "tinker") workModel = current; const model = ctx.modelRegistry.find(target.provider, target.id); if (!model) { sendReport("Work model", `No model registered as \`${target.provider}/${target.id}\`.`, "warning"); return; } if (!await pi.setModel(model)) { sendReport("Work model", `Could not select \`${describeModel(model)}\`. Check that provider's API key.`, "error"); return; } refreshStatus(ctx); sendReport("Work model", `Now using \`${describeModel(model)}\` for execution.`, "success"); return; }
if (current && current.provider !== "tinker") workModel = current; const requested = parseModelRef(rawArgs); const checkpoint = latestCheckpointId(); const target = requested ?? (checkpoint ? { provider: "tinker", id: checkpoint } : { provider: "tinker", id: DEFAULT_FT_MODEL }); const model = ctx.modelRegistry.find(target.provider, target.id); if (!model) { sendReport("Student model", [ `No model registered as \`${target.provider}/${target.id}\`.`, "Machine registers Tinker chat models (Inkling, Qwen 3.8-27B, and `/tinker use` checkpoints). Pass an explicit id:", "```text", `/machine student ${QWEN_38_27B_256K_MODEL}`, "/machine student thinkingmachines/Inkling-Small:peft:262144", "/machine student tinker://.../sampler_weights/final", "```", ].join("\n"), "warning"); return; } if (!await pi.setModel(model)) { sendReport("Student model", `Could not select \`${describeModel(model)}\`. Is \`TINKER_API_KEY\` set?`, "error"); return; } refreshStatus(ctx); sendReport("Student model", [ `Now using \`${describeModel(model)}\`.`, workModel ? `Daily driver remembered as \`${workModel.provider}/${workModel.id}\`. \`/machine work\` jumps back.` : "`/machine work` will remember the next non-Tinker model you use.", ].join("\n"), "success"); }
function setAutoTrace(args: string, ctx: ExtensionCommandContext) { const token = args.trim().toLowerCase(); let autoTrace = readMachineSettings(ctx.cwd).autoTrace; if (token === "on" || token === "true" || token === "1") autoTrace = true; else if (token === "off" || token === "false" || token === "0") autoTrace = false; else if (token === "status") { sendReport("Auto trace", [ autoTrace ? "On for this project." : "Off for this project (default).", `Observed traces: \`${machineDataFile(ctx.cwd, "observed")}\``, ].join("\n")); return; } else if (!token) { autoTrace = !autoTrace; } else { sendReport("Auto trace", "Usage: `/machine auto`, `/machine auto on`, `/machine auto off`, `/machine auto status`.", "warning"); return; } writeMachineSettings(ctx.cwd, { autoTrace }); refreshStatus(ctx); sendReport("Auto trace", autoTrace ? [ `On. Machine will save completed spans as unlabeled observations in \`${machineDataFile(ctx.cwd, "observed")}\`.`, "Observed trajectories include tool output. Machine does not select them for training by default; inspect them before using or copying them.", ].join("\n") : "Off. Machine will not capture compaction or headless spans automatically.", "success"); }
function showDataStatus(ctx: ExtensionCommandContext) { sendReport("Machine data", formatMachineDataStatus(ctx.cwd)); }
pi.registerCommand("machine", { description: "Parametric learning: auto-trace, fork observed runs, or see /good /bad /tagged /train", handler: async (args, ctx) => { const [subcommand, ...rest] = shellSplit(args); if (!subcommand || subcommand === "help" || subcommand === "--help") { sendReport("machine", [ "```text", "/good [--trace] save last turn", "/bad save rejected turn", "/tagged list saved examples", "/machine auto configure compact + headless observation", "/machine data show the project data store", "/machine fork [note] retrigger the latest observed headless run", "/machine correction --id <id> prompt", "/machine work switch back to the daily driver (Grok, etc.)", "/machine student switch to Inkling, Qwen 3.8-27B, or a Tinker checkpoint", "/train validate|sft|smoke|eval", "```", ].join("\n")); return; } if (subcommand === "auto") { setAutoTrace(rest.join(" "), ctx); return; } if (subcommand === "data") { showDataStatus(ctx); return; } if (subcommand === "fork") { const { positional, options } = parseOptions(rest); const result = await forkObservedRun({ cwd: ctx.cwd, correction: positional.join(" "), spawn: true, inheritStdio: false, id: typeof options.id === "string" ? options.id : undefined, mode: "fork", }); sendReport(result.title, result.body, result.level); return; } if (subcommand === "correction" || subcommand === "correct") { const { positional, options } = parseOptions(rest); const id = typeof options.id === "string" ? options.id : ""; const correction = positional.join(" ").trim(); if (!id || !correction) { sendReport("machine correction", "Usage: `/machine correction --id <id> prompt`", "warning"); return; } const result = await forkObservedRun({ cwd: ctx.cwd, correction, spawn: true, inheritStdio: false, id, mode: "correction", }); sendReport(result.title, result.body, result.level); return; } if (subcommand === "work" || subcommand === "student" || subcommand === "ft" || subcommand === "model") { await switchLane(subcommand === "work" ? "work" : "student", rest.join(" "), ctx); return; } sendReport("machine", `Unknown subcommand: ${subcommand}. Run \`/machine help\`.`, "warning"); }, });
async function saveAutoTrace(options: { cwd: string; entries: SessionMessageLike[]; sessionFile?: string; kind: "auto-compact" | "auto-headless"; trigger: string; note: string; }) { const records = recordsFromSessionEntries(options.entries).filter((record) => record.role !== "reasoning"); if (records.length === 0) return undefined; const outFile = resolveTaggedExamplePath(options.cwd, "observed", undefined, true); return appendTaggedExample({ tag: "observed", messages: trajectoryToMessages(records), records, note: options.note, outFile, sessionFile: options.sessionFile, userId: options.entries.find((entry) => entry.message?.role === "user")?.id, assistantId: [...options.entries].reverse().find((entry) => entry.message?.role === "assistant")?.id, kind: options.kind, trigger: options.trigger, }); }
pi.on("session_before_compact", async (event, ctx) => { if (!readMachineSettings(ctx.cwd).autoTrace) return; const firstKept = event.preparation.firstKeptEntryId; const branch = event.branchEntries as SessionMessageLike[]; const cut = firstKept ? branch.findIndex((entry) => entry.id === firstKept) : -1; const summarized = (cut >= 0 ? branch.slice(0, cut) : branch).filter((entry) => entry.type === "message" && entry.message); try { const saved = await saveAutoTrace({ cwd: ctx.cwd, entries: summarized, sessionFile: ctx.sessionManager.getSessionFile(), kind: "auto-compact", trigger: event.reason, note: `auto compact (${event.reason})`, }); if (!saved) { sendReport("Auto trace", "Compaction fired, but there was no message span to save.", "warning"); return; } refreshStatus(ctx); const records = recordsFromSessionEntries(summarized).filter((record) => record.role !== "reasoning"); const user = records.find((record) => record.role === "user"); sendReport("Auto-saved compact observation", [ `Compaction (${event.reason}) wrote \`${rel(ctx.cwd, resolveTaggedExamplePath(ctx.cwd, "observed", undefined, true))}\` (${saved.count} row${saved.count === 1 ? "" : "s"}).`, saved.row.id ? `id: \`${String(saved.row.id)}\`` : "", user ? `First user: ${previewText(user.content, 200)}` : "", "", summarizeTrajectory(records), ].filter(Boolean).join("\n"), "success"); } catch (error) { sendReport("Auto trace failed", error instanceof Error ? error.message : String(error), "error"); } });
pi.on("session_shutdown", async (event, ctx) => { if (!readMachineSettings(ctx.cwd).autoTrace) return; if (ctx.mode !== "print" && ctx.mode !== "json") return; if (event.reason !== "quit") return; const entries = lastUserTurnEntries(ctx.sessionManager.buildContextEntries() as SessionMessageLike[]); const userText = extractMessageText(entries.find((entry) => entry.message?.role === "user")?.message?.content); if (!userText || userText.startsWith("/")) return; try { const saved = await saveAutoTrace({ cwd: ctx.cwd, entries, sessionFile: ctx.sessionManager.getSessionFile(), kind: "auto-headless", trigger: ctx.mode, note: `auto headless (${ctx.mode})`, }); if (!saved) return; const idNote = saved.row.id ? ` id=${String(saved.row.id)}` : ""; const line = `machine: wrote headless observation to ${rel(ctx.cwd, resolveTaggedExamplePath(ctx.cwd, "observed", undefined, true))} (${saved.count} rows)${idNote}\n`; process.stderr.write(line); } catch (error) { process.stderr.write(`machine: failed to write headless trace: ${error instanceof Error ? error.message : String(error)}\n`); } });}