Add best-of-k diverse generation on forge primitives
ober
2ae31c2e2b4ff69e8ebd657d332dfabae7465e46
--- a/build-binary.ss +++ b/build-binary.ss @@ -144,6 +144,7 @@ "lib/jcode/guardrails/step-enforcer" "lib/jcode/core/workflow-runner" "lib/jcode/core/verified" + "lib/jcode/core/best-of-k" "lib/jcode/provider/sampling" "lib/jcode/provider/provider" "lib/jcode/proxy/convert" new file mode 100644 --- /dev/null +++ b/src/jcode/core/best-of-k.ss @@ -0,0 +1,52 @@ +;;; jcode best-of-k diverse generation — ATLAS Phase-1 (diverse-gen +12.4pp), +;;; expressed entirely on forge's run-workflow. +;;; +;;; ATLAS's V3 ablation (docs/reports/V3_ABLATION_STUDY.md in that repo) found +;;; diverse candidate generation the single largest win — +12.4pp, larger than +;;; self-verified repair (+7.3pp, shipped here as verified.ss). The mechanism: +;;; sample several candidate trajectories and let an oracle select a good one. +;;; forge's verify-gate already *is* that oracle — a candidate only completes if +;;; its required verify step passed (see verified.ss) — so best-of-k needs no +;;; new engine code: run the workflow up to K times and return the first attempt +;;; that completes. +;;; +;;; * each attempt = one full run-workflow trajectory. +;;; * a fresh responder per attempt (RESPONDER-FACTORY is a thunk) so the +;;; provider re-samples. Diversity comes from temperature>0, which +;;; sampling.ss already sets per model; greedy decoding makes the K attempts +;;; identical and best-of-k a no-op — the documented precondition. +;;; * verify is the selector: run-workflow RAISES on a trajectory that never +;;; clears the gate (verify never passed, or the iteration/error budget ran +;;; out), so a raised attempt is a rejected candidate and a returned value is +;;; an accepted one. +;;; +;;; First-success semantics: a binary verify makes any passing candidate +;;; acceptable, so we stop at the first one (also the cheapest — no compute spent +;;; after a pass). If every attempt fails, the last trajectory's error is +;;; re-raised so the caller sees a real failure rather than a silent miss. + +(export run-best-of-k) + +(import :jcode/core/workflow-runner) + +;; Run WORKFLOW up to K times, each attempt driven by a fresh responder from the +;; thunk RESPONDER-FACTORY, returning the first attempt's terminal value that +;; completes. OPT is the run-workflow options assoc, passed through unchanged to +;; every attempt. If every attempt raises, the last error is re-raised; K<1 is +;; treated as 1. Drop-in for run-workflow when you can afford K samples and have +;; a verify-gate to select among them. +(def (run-best-of-k workflow user-message responder-factory k . opt) + (let ((o (if (pair? opt) (car opt) '())) + (kk (if (< k 1) 1 k))) + (let loop ((attempt 0) (last-err #f)) + (if (>= attempt kk) + (if last-err + (raise last-err) + (error 'best-of-k "no attempt produced a result")) + (let ((outcome + (guard (e [#t (cons 'fail e)]) + (cons 'ok (run-workflow workflow user-message + (responder-factory) o))))) + (if (eq? (car outcome) 'ok) + (cdr outcome) + (loop (+ attempt 1) (cdr outcome)))))))) --- a/src/jcode/ui/cli.ss +++ b/src/jcode/ui/cli.ss @@ -31,6 +31,7 @@ :jcode/core/workflow :jcode/core/workflow-runner :jcode/core/verified + :jcode/core/best-of-k :jcode/core/slot-worker :jcode/proxy/server :jcode/eval/scenario @@ -503,6 +504,71 @@ EXAMPLES: (printf "With the gate, the failing verify is repaired and correct code ships;~n") (printf "without it, the premature done ships unwritten code (accuracy #f).~n")) +;; /forge bestofk — ATLAS Phase-1 diverse generation on forge's run-workflow. +;; Two candidate trajectories for one task: candidate 0 writes buggy code that +;; never clears the verify-gate (run-workflow raises → rejected); candidate 1 is +;; a diverse re-sample that verifies and finishes. best-of-1 sees only the buggy +;; one and fails; best-of-2 lets the verify-gate select the passing trajectory. +(def (forge-bestofk-workflow code) + (make-verified-workflow + "bok" "Edit, verify, done." + (make-tool-def + (make-tool-spec "edit" "Write code." '(("type" . "object"))) + (lambda (args) + (vector-set! code 0 (let ((p (assoc "content" args))) (if p (cdr p) ""))) + "edited") + '()) + (make-tool-def + (make-tool-spec "verify" "Run the build/tests." '(("type" . "object"))) + (make-verify-callable + (lambda (args) + (if (equal? (vector-ref code 0) "correct") + (cons #t "all tests passed") + (cons #f "test_basic FAILED: wrong output")))) + '()) + (make-tool-def + (make-tool-spec "done" "Finish." '(("type" . "object"))) + (lambda (args) "SHIPPED") + '()))) + +;; Factory thunk: successive calls hand back candidate 0, then candidate 1, … +(def (forge-bestofk-factory) + (let ((attempt -1)) + (lambda () + (set! attempt (+ attempt 1)) + (let ((script + (if (= attempt 0) + (list (list (make-wtool-call "edit" '(("content" . "buggy")) "draft")) + (list (make-wtool-call "verify" '() "checking")) + (list (make-wtool-call "edit" '(("content" . "buggy")) "still wrong")) + (list (make-wtool-call "verify" '() "checking")) + (list (make-wtool-call "edit" '(("content" . "buggy")) "still wrong")) + (list (make-wtool-call "verify" '() "checking"))) + (list (list (make-wtool-call "edit" '(("content" . "correct")) "resample")) + (list (make-wtool-call "verify" '() "checking")) + (list (make-wtool-call "done" '(("summary" . "fixed")) "done"))))) + (n 0)) + (lambda (messages tool-specs step) + (let ((r (if (< n (length script)) (list-ref script n) + (make-text-response "stuck")))) + (set! n (+ n 1)) r)))))) + +(def (forge-print-bestofk) + (printf "Best-of-k: ATLAS diverse generation on forge's run-workflow.~n") + (printf " run the workflow up to k times; the verify-gate selects the first~n") + (printf " trajectory that passes. Diversity needs temperature>0 (sampling.ss).~n") + (let ((opts (list (cons 'max-iterations 5) (cons 'max-tool-errors 10)))) + (let ((k1 (guard (e [#t 'rejected]) + (run-best-of-k (forge-bestofk-workflow (vector "UNWRITTEN")) + "implement the feature" (forge-bestofk-factory) 1 opts)))) + (printf " best-of-1 result=~a (only the buggy candidate)~n" k1)) + (let ((k2 (guard (e [#t 'rejected]) + (run-best-of-k (forge-bestofk-workflow (vector "UNWRITTEN")) + "implement the feature" (forge-bestofk-factory) 2 opts)))) + (printf " best-of-2 result=~a (verify selected the passing resample)~n" k2))) + (printf "With k=1 the lone buggy candidate is rejected; k=2 resamples and the~n") + (printf "verify-gate selects correct code. Composes with /forge verify.~n")) + ;; `jcode proxy` — serve the configured provider behind the guardrail proxy. (def (proxy-main args) (let loop ((args args) (port 8080) (bind "127.0.0.1")) @@ -531,7 +597,7 @@ EXAMPLES: (let ((cmd (string-trim (substring input 1 (string-length input))))) (cond ((equal? cmd "help") - (display "\nCommands:\n /help Show this help\n /model [name] Show or set model\n /provider [name] Show or set provider\n /plan Switch to PLAN mode (read-only)\n /build Switch to BUILD mode (read+write)\n /mode Show current mode\n /mcp Toggle MCP tools on/off\n /tools List available tools\n /clear Start a new session\n /sessions List saved sessions\n /compact Show message count\n /undo [N] Revert last N checkpoint(s) (default 1)\n /checkpoints List recent shadow-git checkpoints\n /forge [on|off] Show or toggle forge guardrails\n /forge sampling <off|on|strict> Per-model sampling policy\n /forge workflow Describe + self-test the workflow engine\n /forge proxy Describe + self-test the OpenAI-compatible proxy\n /forge ablation Describe + self-test the eval/ablation harness\n /forge verify Describe + self-test the verify-gate (ATLAS verify+repair)\n /quit Exit\n\nMulti-line: end a line with \\ to continue on the next line.\n\n")) + (display "\nCommands:\n /help Show this help\n /model [name] Show or set model\n /provider [name] Show or set provider\n /plan Switch to PLAN mode (read-only)\n /build Switch to BUILD mode (read+write)\n /mode Show current mode\n /mcp Toggle MCP tools on/off\n /tools List available tools\n /clear Start a new session\n /sessions List saved sessions\n /compact Show message count\n /undo [N] Revert last N checkpoint(s) (default 1)\n /checkpoints List recent shadow-git checkpoints\n /forge [on|off] Show or toggle forge guardrails\n /forge sampling <off|on|strict> Per-model sampling policy\n /forge workflow Describe + self-test the workflow engine\n /forge proxy Describe + self-test the OpenAI-compatible proxy\n /forge ablation Describe + self-test the eval/ablation harness\n /forge verify Describe + self-test the verify-gate (ATLAS verify+repair)\n /forge bestofk Describe + self-test best-of-k diverse-gen (ATLAS Phase-1)\n /quit Exit\n\nMulti-line: end a line with \\ to continue on the next line.\n\n")) ((equal? cmd "model") (printf "Provider: ~a~n" (or (current-provider-override) (config-provider))) (printf "Model: ~a~n" (or (current-model-override) (config-model))) @@ -626,6 +692,8 @@ EXAMPLES: (forge-print-ablation)) ((equal? cmd "forge verify") (forge-print-verify)) + ((or (equal? cmd "forge bestofk") (equal? cmd "forge best-of-k")) + (forge-print-bestofk)) ((or (equal? cmd "forge on") (equal? cmd "forge enforce") (equal? cmd "forge enforce on")) (forge-respond-enforced? #t) --- a/test/run.ss +++ b/test/run.ss @@ -24,6 +24,7 @@ (jcode guardrails step-enforcer) (jcode core workflow-runner) (jcode core verified) + (jcode core best-of-k) (jcode proxy convert) (jcode proxy handler) (jcode core slot-worker) @@ -1383,6 +1384,60 @@ (check! "verify-gate flip reforged 1/1" (cdr (assoc "reforged" rates)) '(1 . 1)) (check! "verify-gate flip no_steps 0/1" (cdr (assoc "no_steps" rates)) '(0 . 1))) +(section "=== best-of-k: ATLAS diverse-gen, verify selects ===") +;; Two candidates for one task. Candidate 0 writes buggy code and never gets it +;; past verify within the iteration budget — run-workflow RAISES, so the +;; candidate is rejected. Candidate 1 (a diverse re-sample) writes correct code, +;; verifies, and finishes. best-of-1 sees only the buggy candidate and fails; +;; best-of-2 resamples and the verify-gate selects the passing trajectory. This +;; is ATLAS's Phase-1 diverse-generation win composed with the verify-gate. +(define (bok-workflow code) + (make-verified-workflow "bok" "edit, verify, done" + (make-tool-def (make-tool-spec "edit" "write code" '(("type" . "object"))) + (lambda (a) + (vector-set! code 0 (let ([p (assoc "content" a)]) (if p (cdr p) ""))) + "edited") + '()) + (make-tool-def (make-tool-spec "verify" "run tests" '(("type" . "object"))) + (make-verify-callable + (lambda (a) + (if (equal? (vector-ref code 0) "correct") + (cons #t "all tests passed") + (cons #f "test_basic FAILED: wrong output")))) + '()) + (make-tool-def (make-tool-spec "done" "finish" '(("type" . "object"))) + (lambda (a) "SHIPPED") '()))) +;; A factory thunk: successive calls return candidate 0, then candidate 1, … +(define (bok-factory) + (let ([attempt -1]) + (lambda () + (set! attempt (+ attempt 1)) + (if (= attempt 0) + (scripted-responder + (list (list (make-wtool-call "edit" '(("content" . "buggy")) #f)) + (list (make-wtool-call "verify" '() #f)) + (list (make-wtool-call "edit" '(("content" . "buggy")) #f)) + (list (make-wtool-call "verify" '() #f)) + (list (make-wtool-call "edit" '(("content" . "buggy")) #f)) + (list (make-wtool-call "verify" '() #f)))) + (scripted-responder + (list (list (make-wtool-call "edit" '(("content" . "correct")) #f)) + (list (make-wtool-call "verify" '() #f)) + (list (make-wtool-call "done" '(("summary" . "fixed")) #f)))))))) +(define bok-opts (list (cons 'max-iterations 5) (cons 'max-tool-errors 10))) +;; best-of-1: only the buggy candidate → never verifies → run-workflow raises. +(check! "best-of-1 fails (buggy candidate, no selection)" + (raises? (lambda () (run-best-of-k (bok-workflow (vector "UNWRITTEN")) + "go" (bok-factory) 1 bok-opts))) #t) +;; best-of-2: candidate 1 verifies → returned as the selected trajectory. +(check! "best-of-2 succeeds (verify selects the passing resample)" + (run-best-of-k (bok-workflow (vector "UNWRITTEN")) "go" (bok-factory) 2 bok-opts) + "SHIPPED") +;; k<1 is clamped to a single attempt. +(check! "best-of-0 clamps to 1 (still fails on the buggy candidate)" + (raises? (lambda () (run-best-of-k (bok-workflow (vector "UNWRITTEN")) + "go" (bok-factory) 0 bok-opts))) #t) + ;; ── Results ─────────────────────────────────────────────────────── (printf "~n~a passed, ~a failed~n" pass-count fail-count)