Something went wrong. Try again.
An agent harness for parametric learning.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181import { randomUUID } from "node:crypto";import type { EventStore } from "./event-store.js";import { cloneMessages, firstNewUserMessage, stableHash } from "./text.js";import type { BranchBackend, CandidateRecord, ForkResult, PreferenceRecord, RunResult, SerializedMessage, Side, StreamUpdate,} from "./types.js";
export interface ForkCallbacks { left?: (update: StreamUpdate) => void; right?: (update: StreamUpdate) => void;}
export class BranchEngine { private history: SerializedMessage[] = [];
constructor( readonly backend: BranchBackend, readonly store: EventStore, ) {}
get canonicalHistory(): SerializedMessage[] { return cloneMessages(this.history); }
restore(messages: SerializedMessage[]): void { this.history = cloneMessages(messages); }
async fork(prompt: string, callbacks: ForkCallbacks = {}): Promise<ForkResult> { const turnId = randomUUID(); const parentHistory = cloneMessages(this.history); const [left, right] = await Promise.all([ this.backend.run({ history: cloneMessages(parentHistory), prompt, branchId: `${turnId}:left`, ...(callbacks.left ? { onUpdate: callbacks.left } : {}), }), this.backend.run({ history: cloneMessages(parentHistory), prompt, branchId: `${turnId}:right`, ...(callbacks.right ? { onUpdate: callbacks.right } : {}), }), ]);
return { turnId, createdAt: new Date().toISOString(), prompt, parentHistory, parentHistoryHash: stableHash(parentHistory), left, right, }; }
async select(fork: ForkResult, side: Side): Promise<PreferenceRecord> { const preferred = fork[side]; const rejected = fork[side === "left" ? "right" : "left"]; this.history = cloneMessages(preferred.messages);
const record = this.baseRecord(fork, [this.candidate(fork.left, "left"), this.candidate(fork.right, "right")], { preferredCandidateId: preferred.branchId, rejectedCandidateIds: [rejected.branchId], selection: side, }); await this.store.append(record); return record; }
async revise( fork: ForkResult, side: Side, feedback: string, onUpdate?: (update: StreamUpdate) => void, ): Promise<PreferenceRecord> { const original = fork[side]; const other = fork[side === "left" ? "right" : "left"]; const revisionPrompt = buildRevisionPrompt(fork.prompt, original.response, feedback); const revision = await this.backend.run({ history: cloneMessages(original.messages), prompt: revisionPrompt, branchId: `${fork.turnId}:revision`, ...(onUpdate ? { onUpdate } : {}), });
const originalUserMessage = firstNewUserMessage(fork.parentHistory.length, original.messages); const revisedAssistantMessage = revision.messages.findLast((message) => message.role === "assistant"); if (!originalUserMessage || !revisedAssistantMessage) { throw new Error("Could not squash the correction into a clean canonical turn."); }
this.history = [ ...cloneMessages(fork.parentHistory), structuredClone(originalUserMessage), structuredClone(revisedAssistantMessage), ];
const candidates = [ this.candidate(fork.left, "left"), this.candidate(fork.right, "right"), this.candidate(revision, "revision"), ]; const record = this.baseRecord(fork, candidates, { preferredCandidateId: revision.branchId, rejectedCandidateIds: [original.branchId, other.branchId], selection: side === "left" ? "revise-left" : "revise-right", correction: { feedback, originalCandidateId: original.branchId, revisionPrompt, }, }); await this.store.append(record); return record; }
private candidate(result: RunResult, source: CandidateRecord["source"]): CandidateRecord { return { id: result.branchId, source, response: result.response, reasoning: result.reasoning, messages: cloneMessages(result.messages), model: structuredClone(result.model), }; }
private baseRecord( fork: ForkResult, candidates: CandidateRecord[], decision: Pick< PreferenceRecord, "preferredCandidateId" | "rejectedCandidateIds" | "selection" | "correction" >, ): PreferenceRecord { return { schema: "branch-preference/v1", id: fork.turnId, createdAt: new Date().toISOString(), backend: this.backend.name, systemPrompt: this.backend.systemPrompt, systemPromptHash: stableHash(this.backend.systemPrompt), parentHistory: cloneMessages(fork.parentHistory), parentHistoryHash: fork.parentHistoryHash, prompt: fork.prompt, candidates, preferredCandidateId: decision.preferredCandidateId, rejectedCandidateIds: decision.rejectedCandidateIds, selection: decision.selection, ...(decision.correction ? { correction: decision.correction } : {}), }; }}
export function buildRevisionPrompt(originalPrompt: string, originalResponse: string, feedback: string): string { return [ "The response you just gave violated an important preference or instruction.", "", "Feedback:", feedback, "", "Original user prompt:", originalPrompt, "", "Response being corrected:", originalResponse, "", "Respond to the original user prompt again as if you had already understood this feedback.", "Do not mention the correction, this feedback, or a previous response. Return only the improved response.", ].join("\n");}