From 9e2d177d9e1331dd0493f457a092c4e993085417 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 12 Jul 2026 01:11:58 +0000 Subject: [PATCH 1/8] fix: close sandbox symlink escape, conformance waiver, and capture redaction gaps - MxcSandbox uploadFiles/downloadFiles now verify realpath containment and refuse symlinked write targets, so links created inside the workspace can no longer read or write host files outside it - train({ conformance: false }) waives the conformance requirement instead of feeding a failed-check result into the promotion gate, which rejected every candidate while the rubric said checks were disabled - capture mapInput/mapOutput/redact results are used as returned; undefined no longer falls back to the raw unredacted value, and default serialization always yields a string Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X2LxfjsASWmdVoFANEURWR --- packages/harness/src/sandbox.ts | 28 ++++++++++++++--- packages/harness/test/harness.test.ts | 26 +++++++++++++++- src/training.ts | 14 +++++---- test/training.test.ts | 44 +++++++++++++++++++++++++++ 4 files changed, 101 insertions(+), 11 deletions(-) diff --git a/packages/harness/src/sandbox.ts b/packages/harness/src/sandbox.ts index 41a66f5..fc98ae2 100644 --- a/packages/harness/src/sandbox.ts +++ b/packages/harness/src/sandbox.ts @@ -1,4 +1,4 @@ -import { mkdir, readFile, writeFile } from "node:fs/promises"; +import { lstat, mkdir, readFile, realpath, writeFile } from "node:fs/promises"; import { dirname, isAbsolute, relative, resolve } from "node:path"; import { @@ -72,9 +72,11 @@ export class MxcSandbox extends BaseSandbox { try { const target = this.#path(path); await mkdir(dirname(target), { recursive: true }); + await this.#assertContained(dirname(target)); + if (await isSymlink(target)) throw new Error("path escapes sandbox workspace"); await writeFile(target, content); return { path, error: null }; - } catch (error) { + } catch { return { path, error: "permission_denied" as const }; } }))); @@ -84,8 +86,10 @@ export class MxcSandbox extends BaseSandbox { return this.#perform("sandbox.download", { sandbox: this.id, paths }, () => Promise.all(paths.map(async (path) => { try { - return { path, content: await readFile(this.#path(path)), error: null }; - } catch (error) { + const target = this.#path(path); + await this.#assertContained(target); + return { path, content: await readFile(target), error: null }; + } catch { return { path, content: null, error: "file_not_found" as const }; } }))); @@ -103,4 +107,20 @@ export class MxcSandbox extends BaseSandbox { if (fromWorkspace.startsWith("..") || isAbsolute(fromWorkspace)) throw new Error("path escapes sandbox workspace"); return target; } + + /** Host file access follows symlinks, so containment must hold after resolving them too. */ + async #assertContained(path: string): Promise { + const fromWorkspace = relative(await realpath(this.#workspace), await realpath(path)); + if (fromWorkspace.startsWith("..") || isAbsolute(fromWorkspace)) { + throw new Error("path escapes sandbox workspace"); + } + } +} + +async function isSymlink(path: string): Promise { + try { + return (await lstat(path)).isSymbolicLink(); + } catch { + return false; + } } diff --git a/packages/harness/test/harness.test.ts b/packages/harness/test/harness.test.ts index 22980fc..6f8f15c 100644 --- a/packages/harness/test/harness.test.ts +++ b/packages/harness/test/harness.test.ts @@ -1,4 +1,4 @@ -import { appendFile, mkdtemp } from "node:fs/promises"; +import { appendFile, mkdtemp, readFile, symlink, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -139,6 +139,30 @@ describe("training harness", () => { .toThrow("outside the writable sandbox"); }); + it("refuses symlinked paths that resolve outside the workspace", async () => { + const workspace = await mkdtemp(join(tmpdir(), "ts-autocode-sandbox-links-")); + const outside = await mkdtemp(join(tmpdir(), "ts-autocode-outside-")); + await writeFile(join(outside, "secret.txt"), "secret", "utf8"); + await symlink(outside, join(workspace, "leak")); + await symlink(join(outside, "secret.txt"), join(workspace, "alias.txt")); + const { bus } = await approvedBus(); + const sandbox = new MxcSandbox({ id: "links", workspace, policy: createHarnessPolicy({ workspace }), bus, role: "student" }); + + expect(await sandbox.downloadFiles(["leak/secret.txt", "alias.txt"])).toEqual([ + { path: "leak/secret.txt", content: null, error: "file_not_found" }, + { path: "alias.txt", content: null, error: "file_not_found" }, + ]); + expect(await sandbox.uploadFiles([ + ["leak/implant.txt", new TextEncoder().encode("x")], + ["alias.txt", new TextEncoder().encode("x")], + ])).toEqual([ + { path: "leak/implant.txt", error: "permission_denied" }, + { path: "alias.txt", error: "permission_denied" }, + ]); + await expect(readFile(join(outside, "implant.txt"))).rejects.toThrow(); + expect(await readFile(join(outside, "secret.txt"), "utf8")).toBe("secret"); + }); + it("creates configurable Deep Agent callbacks for the same run model", async () => { const root = await mkdtemp(join(tmpdir(), "ts-autocode-agents-")); const role = (name: string) => { diff --git a/src/training.ts b/src/training.ts index 92ce400..7a46daf 100644 --- a/src/training.ts +++ b/src/training.ts @@ -243,7 +243,9 @@ class TrainingRuntime implements Training { const decision = await evaluatePromotionGate({ candidate, evaluations: verification.evaluations, - conformance: input.conformance ?? true, + // The engine already validated the candidate; `conformance: false` waives the + // requirement rather than reporting a failed check to the gate. + conformance: true, ...(input.minScore === undefined ? {} : { minScore: input.minScore }), ...(input.minPassRate === undefined ? {} : { minPassRate: input.minPassRate }), ...(input.policy === undefined ? {} : { policy: input.policy }), @@ -270,7 +272,7 @@ class TrainingRuntime implements Training { const decision = await evaluatePromotionGate({ candidate, evaluations: verification.evaluations, - conformance: input.conformance ?? true, + conformance: true, ...(input.minScore === undefined ? {} : { minScore: input.minScore }), ...(input.minPassRate === undefined ? {} : { minPassRate: input.minPassRate }), ...(input.policy === undefined ? {} : { policy: input.policy }), @@ -456,9 +458,9 @@ class TrainingRuntime implements Training { if (this.#settings.capture.enabled === false) return; try { const spanContext = span?.spanContext(); - const input = this.#settings.capture.mapInput?.(args, token) ?? args; + const input = this.#settings.capture.mapInput ? this.#settings.capture.mapInput(args, token) : args; const output = error === undefined - ? this.#settings.capture.mapOutput?.(result, token) ?? result + ? (this.#settings.capture.mapOutput ? this.#settings.capture.mapOutput(result, token) : result) : errorMessage(error); const record: TrainingRecord = { id: this.#settings.idFactory(), @@ -490,7 +492,7 @@ class TrainingRuntime implements Training { } #serialize(value: unknown, field: "input" | "output"): string { - const redacted = this.#settings.capture.redact?.(value, field) ?? value; + const redacted = this.#settings.capture.redact ? this.#settings.capture.redact(value, field) : value; return (this.#settings.capture.serialize ?? defaultSerialize)(redacted); } @@ -534,7 +536,7 @@ function isPromise(value: T): value is T & Promise> { function defaultSerialize(value: unknown): string { if (typeof value === "string") return value; try { - return JSON.stringify(value); + return JSON.stringify(value) ?? String(value); } catch { return String(value); } diff --git a/test/training.test.ts b/test/training.test.ts index f696e78..d4c8bc8 100644 --- a/test/training.test.ts +++ b/test/training.test.ts @@ -60,6 +60,22 @@ describe("trainable method capture", () => { expect(await training.records("Router.route")).toEqual([]); }); + it("lets capture mappers redact values to undefined", async () => { + const training = configureTraining({ + tracing: { enabled: false }, + capture: { mapInput: () => undefined, mapOutput: () => undefined }, + }); + class Router { + route(input: string): string { return input; } + } + applyMethodDecorator(Router, "route", trainable("Router.redacted")); + + expect(new Router().route("secret-input")).toBe("secret-input"); + const [record] = await training.records("Router.redacted"); + expect(record?.succeeded).toBe(true); + expect(JSON.stringify(record)).not.toContain("secret-input"); + }); + it("supports the decorator without external source metadata", async () => { const training = configureTraining({}); class Router { @@ -140,6 +156,34 @@ describe("training execution", () => { expect(await readFile(artifact, "utf8")).toContain('"use training"'); }); + it("waives the conformance requirement instead of rejecting every candidate", async () => { + const directory = await mkdtemp(join(tmpdir(), "ts-autocode-conformance-")); + const artifact = join(directory, "echo.ts"); + await writeFile(artifact, `export function echo(input: string): string { + "use training"; + return input; +}\n`); + const training = configureTraining({ + engine: { id: "conformance-test", optimize: async () => ({ implementation: "return input.toUpperCase();" }) }, + source: { files: [artifact] }, + tracing: { enabled: false }, + }); + + const run = await training.train({ + trainable: "echo", + objective: "Uppercase the input", + conformance: false, + evaluation: { + tests: [{ id: "upper", input: "abc", assert: [{ type: "equals", value: "ABC" }] }], + task: (input) => input, + outputDir: join(directory, "agentv"), + }, + }); + + expect(run.outcome).toBe("ready"); + expect(run.final.decision.failures).toEqual([]); + }); + it("requires enough successful runtime traces before evolving code", async () => { const training = configureTraining({ tracing: { enabled: false } }); class Router { From 96e60fa820e7f32d4110acd4c2545e050d9f9ca4 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 12 Jul 2026 01:24:09 +0000 Subject: [PATCH 2/8] fix: verify sandbox ancestor containment before mkdir creates directories A symlinked ancestor inside the workspace could make recursive mkdir create directories outside it before the containment check ran. Uploads now verify the nearest existing ancestor's realpath first, with a nested-path regression test. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X2LxfjsASWmdVoFANEURWR --- packages/harness/src/sandbox.ts | 18 ++++++++++++++++++ packages/harness/test/harness.test.ts | 6 +++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/packages/harness/src/sandbox.ts b/packages/harness/src/sandbox.ts index fc98ae2..e89ad3f 100644 --- a/packages/harness/src/sandbox.ts +++ b/packages/harness/src/sandbox.ts @@ -71,6 +71,9 @@ export class MxcSandbox extends BaseSandbox { Promise.all(files.map(async ([path, content]) => { try { const target = this.#path(path); + await mkdir(this.#workspace, { recursive: true }); + // Preflight before mkdir: a symlinked ancestor would otherwise create directories outside. + await this.#assertContained(await existingAncestor(dirname(target), this.#workspace)); await mkdir(dirname(target), { recursive: true }); await this.#assertContained(dirname(target)); if (await isSymlink(target)) throw new Error("path escapes sandbox workspace"); @@ -124,3 +127,18 @@ async function isSymlink(path: string): Promise { return false; } } + +async function existingAncestor(path: string, root: string): Promise { + let current = path; + while (current !== root && !(await exists(current))) current = dirname(current); + return current; +} + +async function exists(path: string): Promise { + try { + await lstat(path); + return true; + } catch { + return false; + } +} diff --git a/packages/harness/test/harness.test.ts b/packages/harness/test/harness.test.ts index 6f8f15c..c1c266b 100644 --- a/packages/harness/test/harness.test.ts +++ b/packages/harness/test/harness.test.ts @@ -1,4 +1,4 @@ -import { appendFile, mkdtemp, readFile, symlink, writeFile } from "node:fs/promises"; +import { appendFile, mkdtemp, readFile, stat, symlink, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -161,6 +161,10 @@ describe("training harness", () => { ]); await expect(readFile(join(outside, "implant.txt"))).rejects.toThrow(); expect(await readFile(join(outside, "secret.txt"), "utf8")).toBe("secret"); + + expect(await sandbox.uploadFiles([["leak/sub/nested.txt", new TextEncoder().encode("x")]])) + .toEqual([{ path: "leak/sub/nested.txt", error: "permission_denied" }]); + await expect(stat(join(outside, "sub"))).rejects.toThrow(); }); it("creates configurable Deep Agent callbacks for the same run model", async () => { From bdf2c56d0a1a9f14b14655c93048dd856e33c485 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 12 Jul 2026 01:24:09 +0000 Subject: [PATCH 3/8] feat!: infer trainable identity and expose a ready-to-use training runtime - @trainable() now infers its identity from the decorated class and method; the optional parameter is a symbol (defineTrainable(...).symbol or Symbol.for(...)), and string identities are rejected - source discovery resolves bare @trainable()/@trainable to the enclosing ClassName.method, token .symbol accesses, and registered Symbol.for ids - the new training export delegates lazily to the configured runtime, so the "use training" directive is the only required consumer code and configureTraining() is purely an optional settings override Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X2LxfjsASWmdVoFANEURWR --- README.md | 45 ++++++++++++++++++++++++------------- docs/architecture.md | 10 +++++---- src/index.ts | 1 + src/source.ts | 52 +++++++++++++++++++++++++++++-------------- src/token.ts | 11 +++++++++ src/training.ts | 43 +++++++++++++++++++++++++++++++---- test/source.test.ts | 18 ++++++++++++++- test/training.test.ts | 28 ++++++++++++++++++----- 8 files changed, 162 insertions(+), 46 deletions(-) diff --git a/README.md b/README.md index fd307bd..4f6e8a1 100644 --- a/README.md +++ b/README.md @@ -31,10 +31,6 @@ Place the literal directive first in a function or method body. No import, decorator, wrapper, registration call, or source-region argument is required. ```ts -import { configureTraining } from "ts-autocode"; - -const training = configureTraining(); - class Router { route(input: string): string { "use training"; @@ -57,17 +53,33 @@ unchanged because the directive is the marker; there is no runtime proxy. ## Optional runtime capture The decorator is optional when calls must be intercepted for runtime capture. -It accepts only the trainable identity; global configuration controls capture -and tracing. The decorated method is the source target, so callers never provide -source metadata. +Identity is inferred from the decorated class and method, so nothing is +declared twice; global configuration controls capture and tracing. The +decorated method is the source target, so callers never provide source +metadata. + +```ts +import { trainable } from "ts-autocode"; + +class Router { + @trainable() + route(input: string): string { + return input; + } +} +``` + +The inferred identity above is `Router.route`. Passing an identity is optional +and takes a symbol, for callers that need a durable id detached from the class +name: ```ts import { defineTrainable, trainable } from "ts-autocode"; -const route = defineTrainable("Router.route"); +const route = defineTrainable("legacy.route"); class Router { - @trainable(route) + @trainable(route.symbol) route(input: string): string { return input; } @@ -76,14 +88,16 @@ class Router { A token contains a durable string id and stable `Symbol.for(...)` symbol. The same id binds the method, captures, AgentV results, optimizer candidate, and -promotion decision. String identities such as `@trainable("Router.route")` are -also accepted. +promotion decision. ## Train and promote -AgentV owns eval definitions, graders, traces, scores, and result types. +AgentV owns eval definitions, graders, traces, scores, and result types. The +`training` export is ready to use without any setup call. ```ts +import { training } from "ts-autocode"; + const tests = [ { id: "billing", @@ -171,9 +185,10 @@ AgentV's `workers` option parallelizes live-trace and candidate evals. Independe trainables can be evolved concurrently by the application, while the configured engine and store remain injectable. -`configureTraining(settings)` is the single public runtime configuration entry -point. Settings are optional. The default Ax implementation reads -`OPENAI_API_KEY` from the configured secret provider or process environment. +Configuration is optional: the exported `training` runtime works out of the +box, and `configureTraining(settings)` only overrides its settings. The default +Ax implementation reads `OPENAI_API_KEY` from the configured secret provider or +process environment. Provider-specific Ax tuning remains isolated to the optional `ts-autocode/ax` adapter and is passed through the provider-neutral `engine` slot. diff --git a/docs/architecture.md b/docs/architecture.md index 293e359..d11a415 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -14,10 +14,12 @@ used. ## Runtime capture -The optional `@trainable(id)` decorator intercepts calls without accepting -capture or tracing options. Global `configureTraining()` settings determine -whether calls are captured or traced and how values are serialized and -redacted. The target is always the decorated method. Calls preserve `this`, +The optional `@trainable()` decorator intercepts calls without accepting +capture or tracing options. Identity is inferred from the decorated class and +method; an explicit symbol identity is optional. The exported `training` +runtime works without configuration, and global `configureTraining()` settings +determine whether calls are captured or traced and how values are serialized +and redacted. The target is always the decorated method. Calls preserve `this`, arguments, synchronous or asynchronous return behavior, and thrown errors. Captured traces use AgentV's `Trace`; spans use official OpenTelemetry and diff --git a/src/index.ts b/src/index.ts index 4575735..e31c087 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,6 +1,7 @@ export { configureTraining, trainable, + training, } from "./training.js"; export type { CaptureSettings, diff --git a/src/source.ts b/src/source.ts index 98afefd..4697b86 100644 --- a/src/source.ts +++ b/src/source.ts @@ -3,7 +3,7 @@ import { dirname, extname, resolve } from "node:path"; import ts from "typescript"; import { digest } from "./canonical.js"; -import type { TrainableId } from "./token.js"; +import { trainableIdFromKey, type TrainableId } from "./token.js"; const trainingDirective = "use training"; @@ -83,8 +83,9 @@ function discoverSourceFile( if (ts.isMethodDeclaration(node) && node.body && node.name) { const methodName = propertyName(node.name, sourceFile); const directive = firstDirective(node.body); - const decoratorId = trainableDecoratorId(node, sourceFile, tokenIds); - if (directive || decoratorId) { + const decorator = trainableDecorator(node); + if (directive || decorator) { + const decoratorId = decorator && trainableDecoratorId(decorator, sourceFile, tokenIds); targets.push(targetFor(node, sourceFile, artifactRef, decoratorId ?? `${className ?? "Anonymous"}.${methodName}`, className)); } return; @@ -149,26 +150,43 @@ function firstDirective(body: ts.Block): ts.ExpressionStatement | undefined { : undefined; } +function trainableDecorator(node: ts.MethodDeclaration): ts.Decorator | undefined { + return ts.getDecorators(node)?.find((item) => { + const expression = ts.isCallExpression(item.expression) ? item.expression.expression : item.expression; + return ts.isIdentifier(expression) && expression.text === "trainable"; + }); +} + +/** Resolves an explicit decorator identity; undefined means infer from the method. */ function trainableDecoratorId( - node: ts.MethodDeclaration, + decorator: ts.Decorator, sourceFile: ts.SourceFile, tokens: ReadonlyMap, ): string | undefined { - const decorator = ts.getDecorators(node)?.find((item) => { - const expression = ts.isCallExpression(item.expression) ? item.expression.expression : item.expression; - return ts.isIdentifier(expression) && expression.text === "trainable"; - }); - if (!decorator) return undefined; - if (!ts.isCallExpression(decorator.expression) || decorator.expression.arguments.length === 0) { - throw new TypeError(`@trainable requires a token or id in ${sourceFile.fileName}`); - } + if (!ts.isCallExpression(decorator.expression) || decorator.expression.arguments.length === 0) return undefined; const argument = decorator.expression.arguments[0] as ts.Expression; - if (ts.isStringLiteralLike(argument)) return argument.text; - if (ts.isIdentifier(argument)) { - const resolved = tokens.get(argument.text); - if (resolved) return resolved; + const id = symbolText(argument) ?? + (ts.isPropertyAccessExpression(argument) && ts.isIdentifier(argument.expression) && argument.name.text === "symbol" + ? tokens.get(argument.expression.text) + : ts.isIdentifier(argument) + ? tokens.get(argument.text) + : undefined); + if (id === undefined) { + throw new TypeError( + `@trainable identity must be a symbol (defineTrainable(...).symbol or Symbol.for(...)) or omitted to infer in ${sourceFile.fileName}`, + ); } - throw new TypeError(`@trainable identity must be a string or defineTrainable token in ${sourceFile.fileName}`); + return id; +} + +function symbolText(expression: ts.Expression): string | undefined { + if (!ts.isCallExpression(expression)) return undefined; + const callee = expression.expression; + const isSymbol = (ts.isIdentifier(callee) && callee.text === "Symbol") || + (ts.isPropertyAccessExpression(callee) && ts.isIdentifier(callee.expression) && + callee.expression.text === "Symbol" && callee.name.text === "for"); + const value = expression.arguments[0]; + return isSymbol && value && ts.isStringLiteralLike(value) ? trainableIdFromKey(value.text) : undefined; } function tokenDeclarations(sourceFile: ts.SourceFile): ReadonlyMap { diff --git a/src/token.ts b/src/token.ts index 53ae69c..5478a06 100644 --- a/src/token.ts +++ b/src/token.ts @@ -27,3 +27,14 @@ export function defineTrainable(id: string): TrainableToken { export function toTrainableToken(identity: TrainableIdentity): TrainableToken { return typeof identity === "string" ? defineTrainable(identity) : identity; } + +/** Strips the library prefix so registered symbols and raw ids share one durable id space. */ +export function trainableIdFromKey(key: string): string { + return key.startsWith(`${tokenPrefix}:`) ? key.slice(tokenPrefix.length + 1) : key; +} + +export function trainableTokenFromSymbol(identity: symbol): TrainableToken { + const key = Symbol.keyFor(identity) ?? identity.description ?? ""; + if (!key.trim()) throw new TypeError("trainable symbol must carry a registry key or description"); + return defineTrainable(trainableIdFromKey(key)); +} diff --git a/src/training.ts b/src/training.ts index 7a46daf..6f36c7a 100644 --- a/src/training.ts +++ b/src/training.ts @@ -32,7 +32,13 @@ import { type SourceSettings, type TrainableTarget, } from "./source.js"; -import { defineTrainable, toTrainableToken, type TrainableIdentity, type TrainableToken } from "./token.js"; +import { + defineTrainable, + toTrainableToken, + trainableTokenFromSymbol, + type TrainableIdentity, + type TrainableToken, +} from "./token.js"; const trainableAttribute = "ts_autocode.trainable.id"; @@ -511,20 +517,49 @@ export function configureTraining(settings: TrainingSettings = {}): Training { return configuredTraining; } -/** Decorator form: `@trainable("Router.route")`. */ -export function trainable(identity: TrainableIdentity): TrainableDecorator { - const token = toTrainableToken(identity); +/** Default runtime: the "use training" directive is the only required marker. + * `configureTraining()` is optional and only overrides settings; each call + * delegates to the current runtime so later configuration still applies. */ +export const training: Training = Object.freeze({ + records: (identity) => runtime().records(identity), + evaluate: (identity, config) => runtime().evaluate(identity, config), + evaluateCandidate: (candidate, config) => runtime().evaluateCandidate(candidate, config), + train: (input) => runtime().train(input), + evolve: (input) => runtime().evolve(input), + optimize: (input) => runtime().optimize(input), + optimizeAll: (inputs) => runtime().optimizeAll(inputs), + promote: (candidate, decision) => runtime().promote(candidate, decision), + revert: (snapshot) => runtime().revert(snapshot), + flush: () => runtime().flush(), +}); + +/** Decorator form: `@trainable()`. Identity is inferred from the decorated class and + * method; pass a symbol (for example `defineTrainable("Router.route").symbol`) only + * to override the inferred id. */ +export function trainable(identity?: symbol): TrainableDecorator { + if (identity !== undefined && typeof identity !== "symbol") { + throw new TypeError("trainable identity must be a symbol; omit it to infer from the decorated method"); + } + const explicit = identity === undefined ? undefined : trainableTokenFromSymbol(identity); return function ( method: (this: This, ...args: Args) => Result, context: ClassMethodDecoratorContext Result>, ) { const name = String(context.name); + let token = explicit; return function (this: This, ...args: Args): Result { + token ??= defineTrainable(`${inferredClassName(this) ?? "Anonymous"}.${name}`); return runtime().invoke(this, method, args, token, name); }; }; } +function inferredClassName(thisValue: unknown): string | undefined { + if (typeof thisValue === "function") return thisValue.name || undefined; + const constructor = (thisValue as { constructor?: unknown } | undefined)?.constructor; + return typeof constructor === "function" && constructor.name ? constructor.name : undefined; +} + function runtime(): TrainingRuntime { return configuredTraining ??= new TrainingRuntime({}); } diff --git a/test/source.test.ts b/test/source.test.ts index 3d7387e..07fd7e9 100644 --- a/test/source.test.ts +++ b/test/source.test.ts @@ -31,7 +31,7 @@ describe("TypeScript trainable discovery", () => { it("resolves decorator tokens without external source metadata", () => { const source = `const route = defineTrainable("router.route"); class Router { - @trainable(route) + @trainable(route.symbol) route(input: string): string { return input; } }`; const [target] = discoverInSource(source, "src/router.ts"); @@ -40,6 +40,22 @@ class Router { expect(target?.implementation).toBe("return input;"); }); + it("infers the id from the decorated class and method when the decorator has no argument", () => { + const source = `class Router { + @trainable() + route(input: string): string { return input; } +}`; + expect(discoverInSource(source, "src/router.ts")[0]?.id).toBe("Router.route"); + }); + + it("resolves registered symbol identities and strips the library prefix", () => { + const source = `class Router { + @trainable(Symbol.for("ts-autocode.trainable:custom.route")) + route(input: string): string { return input; } +}`; + expect(discoverInSource(source, "src/router.ts")[0]?.id).toBe("custom.route"); + }); + it("resolves imported trainable tokens through the TypeScript program", async () => { const output = "test/output/source"; await mkdir(output, { recursive: true }); diff --git a/test/training.test.ts b/test/training.test.ts index d4c8bc8..a11884b 100644 --- a/test/training.test.ts +++ b/test/training.test.ts @@ -9,6 +9,7 @@ import { configureTraining, defineTrainable, trainable, + training as defaultTraining, type TrainingEngine, } from "../src/index.js"; import { discoverInSource } from "../src/source.js"; @@ -41,6 +42,23 @@ describe("trainable identity", () => { const second = defineTrainable("Router.route"); expect(first.symbol).toBe(second.symbol); }); + + it("infers the decorator identity from the decorated class and method", async () => { + configureTraining({ tracing: { enabled: false } }); + class InferredRouter { + route(input: string): string { return input; } + } + applyMethodDecorator(InferredRouter, "route", trainable()); + + expect(new InferredRouter().route("billing")).toBe("billing"); + const [record] = await defaultTraining.records("InferredRouter.route"); + expect(record?.trainableId).toBe("InferredRouter.route"); + expect(record?.succeeded).toBe(true); + }); + + it("rejects non-symbol decorator identities", () => { + expect(() => trainable("Router.route" as never)).toThrow("must be a symbol"); + }); }); describe("trainable method capture", () => { @@ -53,7 +71,7 @@ describe("trainable method capture", () => { class Router { route(input: string): string { return input; } } - applyMethodDecorator(Router, "route", trainable("Router.route")); + applyMethodDecorator(Router, "route", trainable()); expect(new Router().route("billing")).toBe("billing"); expect(startActiveSpan).not.toHaveBeenCalled(); @@ -68,7 +86,7 @@ describe("trainable method capture", () => { class Router { route(input: string): string { return input; } } - applyMethodDecorator(Router, "route", trainable("Router.redacted")); + applyMethodDecorator(Router, "route", trainable(defineTrainable("Router.redacted").symbol)); expect(new Router().route("secret-input")).toBe("secret-input"); const [record] = await training.records("Router.redacted"); @@ -83,7 +101,7 @@ describe("trainable method capture", () => { throw new Error("boom"); } } - applyMethodDecorator(Router, "fail", trainable("Router.fail")); + applyMethodDecorator(Router, "fail", trainable()); await expect(new Router().fail()).rejects.toThrow("boom"); const [record] = await training.records("Router.fail"); @@ -136,7 +154,7 @@ describe("training execution", () => { class RuntimeNormalizer { normalize(input: string): string { return input.toUpperCase(); } } - applyMethodDecorator(RuntimeNormalizer, "normalize", trainable("liveNormalize")); + applyMethodDecorator(RuntimeNormalizer, "normalize", trainable(defineTrainable("liveNormalize").symbol)); const normalize = new RuntimeNormalizer(); normalize.normalize("alpha"); normalize.normalize("beta"); @@ -189,7 +207,7 @@ describe("training execution", () => { class Router { route(input: string): string { return input; } } - applyMethodDecorator(Router, "route", trainable("Router.live")); + applyMethodDecorator(Router, "route", trainable(defineTrainable("Router.live").symbol)); new Router().route("one"); await expect(training.evolve({ From aa14ac1f421903e46c117f0b3b623cd0465316a2 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 12 Jul 2026 01:40:49 +0000 Subject: [PATCH 4/8] refactor!: extract provider-neutral runtime into ts-autocode-training Move tokens, records, engine contracts, evaluation, promotion, source discovery, and the training runtime into an independent ts-autocode-training workspace package. The package has no optimizer or execution provider: TrainingEngine and the new ImplementationExecutor are injected boundaries, with provideTrainingDefaults() registering lazy fallbacks. The root ts-autocode package keeps the Ax provider and AxJSRuntime executor where they were, wires them in as defaults, and re-exports the same public API. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X2LxfjsASWmdVoFANEURWR --- README.md | 6 +- package-lock.json | 36 +++++----- package.json | 12 ++-- packages/training/README.md | 29 ++++++++ packages/training/package.json | 48 +++++++++++++ {src => packages/training/src}/canonical.ts | 0 {src => packages/training/src}/engine.ts | 8 +++ {src => packages/training/src}/evaluation.ts | 0 packages/training/src/index.ts | 49 ++++++++++++++ {src => packages/training/src}/promotion.ts | 0 {src => packages/training/src}/records.ts | 0 {src => packages/training/src}/source.ts | 0 {src => packages/training/src}/token.ts | 0 {src => packages/training/src}/training.ts | 42 ++++++++++-- .../training/test}/engine.test.ts | 0 .../training/test}/evaluation.test.ts | 12 +++- .../training/test}/promotion.test.ts | 8 ++- .../training/test}/source.test.ts | 0 .../training/test}/training.test.ts | 6 ++ packages/training/tsconfig.json | 8 +++ packages/training/tsconfig.test.json | 12 ++++ src/execution.ts | 2 +- src/index.ts | 67 ++++++++++--------- src/providers/ax.ts | 3 +- test/ax.test.ts | 2 +- tsconfig.test.json | 2 +- vitest.config.ts | 2 +- 27 files changed, 283 insertions(+), 71 deletions(-) create mode 100644 packages/training/README.md create mode 100644 packages/training/package.json rename {src => packages/training/src}/canonical.ts (100%) rename {src => packages/training/src}/engine.ts (94%) rename {src => packages/training/src}/evaluation.ts (100%) create mode 100644 packages/training/src/index.ts rename {src => packages/training/src}/promotion.ts (100%) rename {src => packages/training/src}/records.ts (100%) rename {src => packages/training/src}/source.ts (100%) rename {src => packages/training/src}/token.ts (100%) rename {src => packages/training/src}/training.ts (94%) rename {test => packages/training/test}/engine.test.ts (100%) rename {test => packages/training/test}/evaluation.test.ts (86%) rename {test => packages/training/test}/promotion.test.ts (87%) rename {test => packages/training/test}/source.test.ts (100%) rename {test => packages/training/test}/training.test.ts (96%) create mode 100644 packages/training/tsconfig.json create mode 100644 packages/training/tsconfig.test.json diff --git a/README.md b/README.md index 4f6e8a1..587a974 100644 --- a/README.md +++ b/README.md @@ -10,8 +10,10 @@ The normal path keeps the code primitives and agent loop in separate packages: ``` Ax is the default student optimizer. AgentV evaluation and the promotion gate -form the teacher. `ts-autocode` delegates their iterative coordination to the -independent `ts-autocode-harness` package. Its single Flue-style callback loop +form the teacher. The provider-neutral runtime lives in the independent +`ts-autocode-training` package (this package re-exports it with Ax wired in as +the default engine and executor), and iterative coordination is delegated to +the independent `ts-autocode-harness` package. Its single Flue-style callback loop supports configurable student, teacher, judge, and adversary Deep Agents, MXC execution, and a write-ahead approval bus. Consumers can supply callbacks from their own agent lifecycle or optimization pipeline without coupling it to this diff --git a/package-lock.json b/package-lock.json index b9ed488..33958a4 100644 --- a/package-lock.json +++ b/package-lock.json @@ -13,11 +13,8 @@ "packages/*" ], "dependencies": { - "@agentv/core": "^4.42.4", - "@arizeai/openinference-semantic-conventions": "^2.5.0", "@ax-llm/ax": "^23.0.0", - "@opentelemetry/api": "^1.9.1", - "ts-autocode-harness": "0.1.0", + "ts-autocode-training": "0.1.0", "typescript": "^5.9.3" }, "devDependencies": { @@ -604,9 +601,6 @@ "cpu": [ "arm64" ], - "libc": [ - "glibc" - ], "license": "SEE LICENSE IN LICENSE.md", "optional": true, "os": [ @@ -623,9 +617,6 @@ "cpu": [ "x64" ], - "libc": [ - "glibc" - ], "license": "SEE LICENSE IN LICENSE.md", "optional": true, "os": [ @@ -642,9 +633,6 @@ "cpu": [ "arm64" ], - "libc": [ - "musl" - ], "license": "SEE LICENSE IN LICENSE.md", "optional": true, "os": [ @@ -661,9 +649,6 @@ "cpu": [ "x64" ], - "libc": [ - "musl" - ], "license": "SEE LICENSE IN LICENSE.md", "optional": true, "os": [ @@ -3281,6 +3266,10 @@ "resolved": "packages/harness", "link": true }, + "node_modules/ts-autocode-training": { + "resolved": "packages/training", + "link": true + }, "node_modules/tslib": { "version": "2.8.1", "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", @@ -3580,6 +3569,21 @@ "engines": { "node": ">=20" } + }, + "packages/training": { + "name": "ts-autocode-training", + "version": "0.1.0", + "license": "MIT", + "dependencies": { + "@agentv/core": "^4.42.4", + "@arizeai/openinference-semantic-conventions": "^2.5.0", + "@opentelemetry/api": "^1.9.1", + "ts-autocode-harness": "0.1.0", + "typescript": "^5.9.3" + }, + "engines": { + "node": ">=20" + } } } } diff --git a/package.json b/package.json index b41d5ab..9f37253 100644 --- a/package.json +++ b/package.json @@ -35,13 +35,14 @@ "url": "https://github.com/Tyler-R-Kendrick/ts-autocode/issues" }, "scripts": { - "build": "npm run build:harness && npm run build:core", + "build": "npm run build:harness && npm run build:training && npm run build:core", "build:core": "node -e \"require('node:fs').rmSync('dist', { recursive: true, force: true })\" && tsc -p tsconfig.json", "build:harness": "node -e \"require('node:fs').rmSync('packages/harness/dist', { recursive: true, force: true })\" && tsc -p packages/harness/tsconfig.json", - "typecheck": "npm run build:harness && tsc --noEmit -p tsconfig.test.json && tsc --noEmit -p packages/harness/tsconfig.test.json", + "typecheck": "npm run build:harness && npm run build:training && tsc --noEmit -p tsconfig.test.json && tsc --noEmit -p packages/harness/tsconfig.test.json && tsc --noEmit -p packages/training/tsconfig.test.json", "test": "node test/run.mjs", "check": "npm run typecheck && npm test && npm run build:core", - "prepublishOnly": "npm run check" + "prepublishOnly": "npm run check", + "build:training": "node -e \"require('node:fs').rmSync('packages/training/dist', { recursive: true, force: true })\" && tsc -p packages/training/tsconfig.json" }, "keywords": [ "agentv", @@ -55,11 +56,8 @@ "vitest": "^4.1.9" }, "dependencies": { - "@agentv/core": "^4.42.4", - "@arizeai/openinference-semantic-conventions": "^2.5.0", "@ax-llm/ax": "^23.0.0", - "@opentelemetry/api": "^1.9.1", - "ts-autocode-harness": "0.1.0", + "ts-autocode-training": "0.1.0", "typescript": "^5.9.3" }, "overrides": { diff --git a/packages/training/README.md b/packages/training/README.md new file mode 100644 index 0000000..0e7489c --- /dev/null +++ b/packages/training/README.md @@ -0,0 +1,29 @@ +# ts-autocode-training + +Provider-neutral runtime for training and safely rewriting directive-marked +TypeScript functions. This package owns discovery of `"use training"` methods, +runtime capture, AgentV evaluation, candidate validation, the promotion gate, +and the bounded student/teacher loop (via `ts-autocode-harness`). + +It has **no optimizer or execution provider**: `TrainingEngine` (candidate +optimization) and `ImplementationExecutor` (running proposed bodies) are +injected boundaries. Supply them per runtime through `TrainingSettings.engine` +and `TrainingSettings.executor`, or register lazy defaults once with +`provideTrainingDefaults(...)` — that is how the `ts-autocode` package wires Ax +as the default engine and Ax's JavaScript sandbox as the default executor. + +```ts +import { configureTraining, provideTrainingDefaults } from "ts-autocode-training"; + +provideTrainingDefaults({ + engine: () => myEngine, + executor: (target, implementation, args) => myRunner.run(target, implementation, args), +}); +``` + +Most applications should depend on [`ts-autocode`](../../README.md), which +re-exports this package's API with Ax defaults already registered. + +## License + +[MIT](../../LICENSE) diff --git a/packages/training/package.json b/packages/training/package.json new file mode 100644 index 0000000..0c2de03 --- /dev/null +++ b/packages/training/package.json @@ -0,0 +1,48 @@ +{ + "name": "ts-autocode-training", + "version": "0.1.0", + "description": "Provider-neutral runtime for training and safely rewriting directive-marked TypeScript functions.", + "repository": { + "type": "git", + "url": "https://github.com/Tyler-R-Kendrick/ts-autocode.git", + "directory": "packages/training" + }, + "homepage": "https://github.com/Tyler-R-Kendrick/ts-autocode/tree/main/packages/training#readme", + "bugs": { + "url": "https://github.com/Tyler-R-Kendrick/ts-autocode/issues" + }, + "license": "MIT", + "type": "module", + "sideEffects": false, + "files": [ + "dist" + ], + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + } + }, + "engines": { + "node": ">=20" + }, + "scripts": { + "build": "tsc -p tsconfig.json", + "prepublishOnly": "npm run build" + }, + "dependencies": { + "@agentv/core": "^4.42.4", + "@arizeai/openinference-semantic-conventions": "^2.5.0", + "@opentelemetry/api": "^1.9.1", + "ts-autocode-harness": "0.1.0", + "typescript": "^5.9.3" + }, + "keywords": [ + "agentv", + "code-optimization", + "training-runtime", + "rewriter" + ] +} diff --git a/src/canonical.ts b/packages/training/src/canonical.ts similarity index 100% rename from src/canonical.ts rename to packages/training/src/canonical.ts diff --git a/src/engine.ts b/packages/training/src/engine.ts similarity index 94% rename from src/engine.ts rename to packages/training/src/engine.ts index 484578f..1dd2739 100644 --- a/src/engine.ts +++ b/packages/training/src/engine.ts @@ -50,6 +50,14 @@ export interface TrainingEngine { optimize(request: OptimizeRequest, context: EngineContext): Promise; } +/** Runs a proposed implementation against arguments in provider-owned isolation. */ +export type ImplementationExecutor = ( + target: TrainableTarget, + implementation: string, + args: readonly unknown[], + options?: Readonly<{ timeoutMs?: number; signal?: AbortSignal }>, +) => Promise; + export async function optimizeCandidate( engine: TrainingEngine, request: OptimizeRequest, diff --git a/src/evaluation.ts b/packages/training/src/evaluation.ts similarity index 100% rename from src/evaluation.ts rename to packages/training/src/evaluation.ts diff --git a/packages/training/src/index.ts b/packages/training/src/index.ts new file mode 100644 index 0000000..f6e301b --- /dev/null +++ b/packages/training/src/index.ts @@ -0,0 +1,49 @@ +export { + configureTraining, + provideTrainingDefaults, + trainable, + training, +} from "./training.js"; +export type { + CaptureSettings, + CandidateEvalConfig, + EvolveInput, + EvolveResult, + OptimizeInput, + TrainInput, + Training, + TrainingProviders, + TrainingRound, + TrainingRun, + TrainingSettings, + TracingSettings, +} from "./training.js"; + +export { defineTrainable } from "./token.js"; +export type { TrainableId, TrainableIdentity, TrainableToken } from "./token.js"; + +export { discoverInSource, discoverTrainables } from "./source.js"; +export type { SourceSettings, TrainableTarget } from "./source.js"; + +export { applyCandidate } from "./engine.js"; +export type { + BoundEvaluation, + CandidatePatch, + EngineCandidate, + EngineContext, + ImplementationExecutor, + OptimizeRequest, + SecretProvider, + TrainingEngine, +} from "./engine.js"; + +export { evaluatePromotionGate, promoteCandidate, revertPromotion } from "./promotion.js"; +export type { + PromotionDecision, + PromotionGateInput, + PromotionResult, + PromotionSnapshot, +} from "./promotion.js"; + +export { createMemoryTrainingStore } from "./records.js"; +export type { TrainingRecord, TrainingStore } from "./records.js"; diff --git a/src/promotion.ts b/packages/training/src/promotion.ts similarity index 100% rename from src/promotion.ts rename to packages/training/src/promotion.ts diff --git a/src/records.ts b/packages/training/src/records.ts similarity index 100% rename from src/records.ts rename to packages/training/src/records.ts diff --git a/src/source.ts b/packages/training/src/source.ts similarity index 100% rename from src/source.ts rename to packages/training/src/source.ts diff --git a/src/token.ts b/packages/training/src/token.ts similarity index 100% rename from src/token.ts rename to packages/training/src/token.ts diff --git a/src/training.ts b/packages/training/src/training.ts similarity index 94% rename from src/training.ts rename to packages/training/src/training.ts index 6f36c7a..4224406 100644 --- a/src/training.ts +++ b/packages/training/src/training.ts @@ -11,11 +11,11 @@ import { optimizeCandidate, type BoundEvaluation, type CandidatePatch, + type ImplementationExecutor, type SecretProvider, type TrainingEngine, } from "./engine.js"; import { evaluateTrainable, type TrainableEvalRun } from "./evaluation.js"; -import { executeImplementation } from "./execution.js"; import { evaluatePromotionGate, promoteCandidate, @@ -24,7 +24,6 @@ import { type PromotionResult, type PromotionSnapshot, } from "./promotion.js"; -import { createAxEngine } from "./providers/ax.js"; import { createMemoryTrainingStore, type TrainingRecord, type TrainingStore } from "./records.js"; import { discoverTrainables, @@ -61,6 +60,7 @@ export interface TracingSettings { export interface TrainingSettings { readonly engine?: TrainingEngine; + readonly executor?: ImplementationExecutor; readonly source?: SourceSettings; readonly store?: TrainingStore; readonly secrets?: SecretProvider; @@ -149,7 +149,7 @@ export type TrainableDecorator = ( class TrainingRuntime implements Training { readonly #settings: Required> & Omit; - readonly #engine: TrainingEngine; + #engine: TrainingEngine | undefined; readonly #store: TrainingStore; readonly #tracer: Tracer; readonly #pending = new Set>(); @@ -170,11 +170,27 @@ class TrainingRuntime implements Training { tracing: settings.tracing ?? {}, variables: Object.freeze({ ...settings.variables }), }; - this.#engine = settings.engine ?? createAxEngine(); this.#store = settings.store ?? createMemoryTrainingStore(); this.#tracer = this.#settings.tracing.tracer ?? trace.getTracer("ts-autocode"); } + #engineFor(override?: TrainingEngine): TrainingEngine { + if (override) return override; + this.#engine ??= this.#settings.engine ?? defaultProviders.engine?.(); + if (!this.#engine) { + throw new Error('no training engine is configured; import "ts-autocode" for the Ax default or set TrainingSettings.engine'); + } + return this.#engine; + } + + #executorOrThrow(): ImplementationExecutor { + const executor = this.#settings.executor ?? defaultProviders.executor; + if (!executor) { + throw new Error('candidate execution requires an executor; import "ts-autocode" or set TrainingSettings.executor'); + } + return executor; + } + async records(identity?: TrainableIdentity): Promise { await this.flush(); return this.#store.list(identity === undefined ? undefined : toTrainableToken(identity).id); @@ -189,12 +205,13 @@ class TrainingRuntime implements Training { async evaluateCandidate(candidate: CandidatePatch, config: CandidateEvalConfig): Promise { const token = defineTrainable(candidate.trainableId); + const execute = this.#executorOrThrow(); const { signal, ...evaluation } = config; signal?.throwIfAborted(); const evaluated = await evaluateTrainable(token, { ...evaluation, task: async (input) => { - const output = await executeImplementation( + const output = await execute( candidate.target, candidate.implementation, evaluationArgs(input), @@ -343,7 +360,7 @@ class TrainingRuntime implements Training { const target = input.target ?? findTrainable(token.id, this.#settings.source); const records = await this.records(token); return optimizeCandidate( - input.engine ?? this.#engine, + this.#engineFor(input.engine), { trainableId: token.id, objective: input.objective, @@ -517,6 +534,19 @@ export function configureTraining(settings: TrainingSettings = {}): Training { return configuredTraining; } +export interface TrainingProviders { + readonly engine?: () => TrainingEngine; + readonly executor?: ImplementationExecutor; +} + +let defaultProviders: TrainingProviders = {}; + +/** Provider packages call this to supply lazy fallbacks (ts-autocode wires Ax) + * without this package depending on any provider. Explicit settings win. */ +export function provideTrainingDefaults(providers: TrainingProviders): void { + defaultProviders = { ...defaultProviders, ...providers }; +} + /** Default runtime: the "use training" directive is the only required marker. * `configureTraining()` is optional and only overrides settings; each call * delegates to the current runtime so later configuration still applies. */ diff --git a/test/engine.test.ts b/packages/training/test/engine.test.ts similarity index 100% rename from test/engine.test.ts rename to packages/training/test/engine.test.ts diff --git a/test/evaluation.test.ts b/packages/training/test/evaluation.test.ts similarity index 86% rename from test/evaluation.test.ts rename to packages/training/test/evaluation.test.ts index b82998c..4b2da52 100644 --- a/test/evaluation.test.ts +++ b/packages/training/test/evaluation.test.ts @@ -14,6 +14,12 @@ function pipelineTarget(input: string): string { return input; } +const functionExecutor = async ( + target: { readonly parameters: readonly { readonly name: string }[] }, + implementation: string, + args: readonly unknown[], +) => new Function(...target.parameters.map((parameter) => parameter.name), implementation)(...args) as unknown; + describe("AgentV evaluation", () => { it("binds AgentV results to the trainable token", async () => { const token = defineTrainable("Router.route"); @@ -47,7 +53,7 @@ describe("AgentV evaluation", () => { return { implementation: "return input;" }; }, }; - const training = configureTraining({ engine }); + const training = configureTraining({ engine, executor: functionExecutor }); await training.evaluate("Router.route", { tests: [{ id: "identity", input: "hello", assert: [{ type: "equals", value: "hello" }] }], task: (input) => input, @@ -70,7 +76,7 @@ describe("AgentV evaluation", () => { target, implementation: "return input.toUpperCase();", }; - const evaluated = await configureTraining({}).evaluateCandidate(candidate, { + const evaluated = await configureTraining({ executor: functionExecutor }).evaluateCandidate(candidate, { tests: [{ id: "uppercase", input: "hello", assert: [{ type: "equals", value: "HELLO" }] }], outputDir: "test/output/agentv-candidate", }); @@ -92,7 +98,7 @@ describe("AgentV evaluation", () => { return { implementation: round === 1 ? "return input;" : "return input.toUpperCase();" }; }, }; - const training = configureTraining({ engine, source: { files: [import.meta.filename] } }); + const training = configureTraining({ engine, executor: functionExecutor, source: { files: [import.meta.filename] } }); const evaluateCandidate = vi.spyOn(training, "evaluateCandidate"); const run = await training.train({ trainable: "pipelineTarget", diff --git a/test/promotion.test.ts b/packages/training/test/promotion.test.ts similarity index 87% rename from test/promotion.test.ts rename to packages/training/test/promotion.test.ts index 2cf5db9..9f558ca 100644 --- a/test/promotion.test.ts +++ b/packages/training/test/promotion.test.ts @@ -29,10 +29,16 @@ function candidate(): CandidatePatch { }; } +const functionExecutor = async ( + target: { readonly parameters: readonly { readonly name: string }[] }, + implementation: string, + args: readonly unknown[], +) => new Function(...target.parameters.map((parameter) => parameter.name), implementation)(...args) as unknown; + describe("promotion", () => { it("gates with AgentV and reverts only an unchanged promoted method", async () => { const patch = candidate(); - const evaluated = await configureTraining({}).evaluateCandidate(patch, { + const evaluated = await configureTraining({ executor: functionExecutor }).evaluateCandidate(patch, { tests: [{ id: "candidate", input: "route", assert: [{ type: "equals", value: "new" }] }], outputDir: "test/output/agentv-promotion", }); diff --git a/test/source.test.ts b/packages/training/test/source.test.ts similarity index 100% rename from test/source.test.ts rename to packages/training/test/source.test.ts diff --git a/test/training.test.ts b/packages/training/test/training.test.ts similarity index 96% rename from test/training.test.ts rename to packages/training/test/training.test.ts index a11884b..4945d2e 100644 --- a/test/training.test.ts +++ b/packages/training/test/training.test.ts @@ -10,6 +10,7 @@ import { defineTrainable, trainable, training as defaultTraining, + type ImplementationExecutor, type TrainingEngine, } from "../src/index.js"; import { discoverInSource } from "../src/source.js"; @@ -148,6 +149,7 @@ describe("training execution", () => { }); const training = configureTraining({ engine: { id: "live-test", optimize }, + executor: functionExecutor, source: { files: [artifact] }, tracing: { enabled: false }, }); @@ -183,6 +185,7 @@ describe("training execution", () => { }\n`); const training = configureTraining({ engine: { id: "conformance-test", optimize: async () => ({ implementation: "return input.toUpperCase();" }) }, + executor: functionExecutor, source: { files: [artifact] }, tracing: { enabled: false }, }); @@ -218,6 +221,9 @@ describe("training execution", () => { }); }); +const functionExecutor: ImplementationExecutor = async (target, implementation, args) => + new Function(...target.parameters.map((parameter) => parameter.name), implementation)(...args); + function applyMethodDecorator object>( constructor: Class, name: string, diff --git a/packages/training/tsconfig.json b/packages/training/tsconfig.json new file mode 100644 index 0000000..cdc3a24 --- /dev/null +++ b/packages/training/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.json", + "compilerOptions": { + "rootDir": "src", + "outDir": "dist" + }, + "include": ["src"] +} diff --git a/packages/training/tsconfig.test.json b/packages/training/tsconfig.test.json new file mode 100644 index 0000000..250b99d --- /dev/null +++ b/packages/training/tsconfig.test.json @@ -0,0 +1,12 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "noEmit": true, + "rootDir": ".", + "outDir": null, + "declaration": false, + "declarationMap": false, + "sourceMap": false + }, + "include": ["src", "test"] +} diff --git a/src/execution.ts b/src/execution.ts index 7b876c9..10a617e 100644 --- a/src/execution.ts +++ b/src/execution.ts @@ -1,7 +1,7 @@ import { AxJSRuntime } from "@ax-llm/ax"; import ts from "typescript"; -import type { TrainableTarget } from "./source.js"; +import type { TrainableTarget } from "ts-autocode-training"; export async function executeImplementation( target: TrainableTarget, diff --git a/src/index.ts b/src/index.ts index e31c087..9e00965 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,46 +1,51 @@ +import { provideTrainingDefaults } from "ts-autocode-training"; + +import { executeImplementation } from "./execution.js"; +import { createAxEngine } from "./providers/ax.js"; + +provideTrainingDefaults({ engine: () => createAxEngine(), executor: executeImplementation }); + export { + applyCandidate, configureTraining, + createMemoryTrainingStore, + defineTrainable, + discoverTrainables, + evaluatePromotionGate, + promoteCandidate, + revertPromotion, trainable, training, -} from "./training.js"; +} from "ts-autocode-training"; export type { - CaptureSettings, + BoundEvaluation, CandidateEvalConfig, + CandidatePatch, + CaptureSettings, + EngineCandidate, + EngineContext, EvolveInput, EvolveResult, + ImplementationExecutor, OptimizeInput, + OptimizeRequest, + PromotionDecision, + PromotionGateInput, + PromotionResult, + PromotionSnapshot, + SecretProvider, + SourceSettings, TrainInput, + TrainableId, + TrainableIdentity, + TrainableTarget, + TrainableToken, Training, + TrainingEngine, + TrainingRecord, TrainingRound, TrainingRun, TrainingSettings, + TrainingStore, TracingSettings, -} from "./training.js"; - -export { defineTrainable } from "./token.js"; -export type { TrainableId, TrainableIdentity, TrainableToken } from "./token.js"; - -export { discoverTrainables } from "./source.js"; -export type { SourceSettings, TrainableTarget } from "./source.js"; - -export { applyCandidate } from "./engine.js"; -export type { - BoundEvaluation, - CandidatePatch, - EngineCandidate, - EngineContext, - OptimizeRequest, - SecretProvider, - TrainingEngine, -} from "./engine.js"; - -export { evaluatePromotionGate, promoteCandidate, revertPromotion } from "./promotion.js"; -export type { - PromotionDecision, - PromotionGateInput, - PromotionResult, - PromotionSnapshot, -} from "./promotion.js"; - -export { createMemoryTrainingStore } from "./records.js"; -export type { TrainingRecord, TrainingStore } from "./records.js"; +} from "ts-autocode-training"; diff --git a/src/providers/ax.ts b/src/providers/ax.ts index 7282f04..4387309 100644 --- a/src/providers/ax.ts +++ b/src/providers/ax.ts @@ -8,7 +8,8 @@ import { type AxOptimizeOptions, } from "@ax-llm/ax"; -import type { EngineContext, OptimizeRequest, TrainingEngine } from "../engine.js"; +import type { EngineContext, OptimizeRequest, TrainingEngine } from "ts-autocode-training"; + import { executeImplementation } from "../execution.js"; type Service = AxAIService | ((context: EngineContext) => AxAIService | Promise); diff --git a/test/ax.test.ts b/test/ax.test.ts index 1448010..916a487 100644 --- a/test/ax.test.ts +++ b/test/ax.test.ts @@ -2,7 +2,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { configureTraining, defineTrainable, type BoundEvaluation } from "../src/index.js"; import { createAxEngine } from "../src/providers/ax.js"; -import { discoverInSource } from "../src/source.js"; +import { discoverInSource } from "ts-autocode-training"; const mocks = vi.hoisted(() => ({ ax: vi.fn(), diff --git a/tsconfig.test.json b/tsconfig.test.json index 13ccae3..005eb88 100644 --- a/tsconfig.test.json +++ b/tsconfig.test.json @@ -8,6 +8,6 @@ "declarationMap": false, "sourceMap": false }, - "include": ["src", "test", "examples", "packages/harness/src", "packages/harness/test", "vitest.config.ts"], + "include": ["src", "test", "examples", "vitest.config.ts"], "exclude": ["test/output", "test/.agentv", "examples/output"] } diff --git a/vitest.config.ts b/vitest.config.ts index 7f3b5d4..0bda23a 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -2,6 +2,6 @@ import { defineConfig } from "vitest/config"; export default defineConfig({ test: { - include: ["test/**/*.test.ts", "packages/harness/test/**/*.test.ts"], + include: ["test/**/*.test.ts", "packages/harness/test/**/*.test.ts", "packages/training/test/**/*.test.ts"], }, }); From 2fd00363b602344dc55274b0c8487f767a9724ae Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 12 Jul 2026 01:48:57 +0000 Subject: [PATCH 5/8] feat: zero-config evolution via ts-autocode/register runtime patch - node --import ts-autocode/register installs a module load hook that appends guarded instrumentation to every application module containing a "use training" directive, wiring classes and function declarations into the same capture path as the decorator with no consumer code - the training runtime gains a background evolution scheduler: after evolution.minTraces successful captures a trainable is trained, verified, gated, and its source rewritten only when the gate passes; failures surface through onError("evolve") and never block application calls - register enables evolution by default (TS_AUTOCODE_EVOLVE=off or evolution.enabled false to capture only); the base training package leaves it disabled - new wrapTrainable/instrumentTrainable primitives are idempotent and shared with the decorator via a wrapped marker Verified end-to-end: vitest auto-evolve loop rewrites a real artifact, and a live node --import smoke run captures directive-marked class and function calls through the built dist. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X2LxfjsASWmdVoFANEURWR --- README.md | 29 ++++++-- docs/architecture.md | 15 +++++ package.json | 4 ++ packages/training/src/index.ts | 3 + packages/training/src/training.ts | 88 ++++++++++++++++++++++++- packages/training/test/training.test.ts | 54 +++++++++++++++ src/register.ts | 30 +++++++++ src/register/hook.ts | 28 ++++++++ test/register.test.ts | 31 +++++++++ 9 files changed, 275 insertions(+), 7 deletions(-) create mode 100644 src/register.ts create mode 100644 src/register/hook.ts create mode 100644 test/register.test.ts diff --git a/README.md b/README.md index 587a974..8128198 100644 --- a/README.md +++ b/README.md @@ -130,12 +130,31 @@ const promoted = await training.promote(run.final.candidate, run.final.decision) await training.revert(promoted.snapshot); ``` -## Evolve from live traces +## Zero-config evolution -When runtime capture is enabled, `evolve()` turns successful captured calls into -AgentV equality evals, trains a replacement, verifies the candidate against the -same cases, applies the promotion gate, and updates the marked TypeScript body. -The write is explicit: capturing traffic alone never changes source code. +Load the runtime patch once and directive-marked functions evolve from live +traffic with no further code — capture, training, verification, gating, and the +guarded source rewrite all apply automatically: + +```bash +node --import ts-autocode/register ./dist/server.js +``` + +The register hook instruments every `"use training"` function at module load. +Once a trainable accumulates `evolution.minTraces` successful traces (default +3), it is trained against those traces, verified candidate-bound, gated, and — +only when the gate passes — its source body is rewritten. Failures surface +through `TrainingSettings.onError` with the `"evolve"` phase and never block or +alter application calls. Set `TS_AUTOCODE_EVOLVE=off` (or configure +`evolution: { enabled: false }`) to capture without rewriting, and use +`evolution.onEvolved` to observe applied rewrites. + +## Evolve from live traces explicitly + +Without the register patch, `evolve()` is the explicit form of the same loop: +it turns successful captured calls into AgentV equality evals, trains a +replacement, verifies the candidate against the same cases, applies the +promotion gate, and updates the marked TypeScript body. ```ts const result = await training.evolve({ diff --git a/docs/architecture.md b/docs/architecture.md index d11a415..6ba7a03 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -25,6 +25,21 @@ arguments, synchronous or asynchronous return behavior, and thrown errors. Captured traces use AgentV's `Trace`; spans use official OpenTelemetry and OpenInference APIs. +## Zero-config runtime patch + +`ts-autocode/register` installs a `node:module` load hook that appends guarded +instrumentation to every application module containing a `"use training"` +directive, wiring each discovered class method or function declaration into the +same capture path as the decorator. It also enables background evolution by +default: after `evolution.minTraces` successful captures, the runtime runs the +same `evolve()` pipeline — replay evals, candidate verification, promotion +gate, guarded rewrite — off the hot path, reporting failures through +`onError("evolve")`. Calls made during a module's own top-level evaluation +precede its instrumentation; traffic after startup is captured. The training +runtime itself lives in the provider-neutral `ts-autocode-training` package; +`ts-autocode` wires Ax as the default engine and executor via +`provideTrainingDefaults`. + `evolve()` is the explicit runtime-to-source bridge. It converts distinct, successful captured inputs and outputs into official AgentV eval cases, replays them as the baseline, and evaluates generated TypeScript against those same diff --git a/package.json b/package.json index 9f37253..55b5201 100644 --- a/package.json +++ b/package.json @@ -28,6 +28,10 @@ "./ax": { "types": "./dist/providers/ax.d.ts", "import": "./dist/providers/ax.js" + }, + "./register": { + "types": "./dist/register.d.ts", + "import": "./dist/register.js" } }, "homepage": "https://github.com/Tyler-R-Kendrick/ts-autocode#readme", diff --git a/packages/training/src/index.ts b/packages/training/src/index.ts index f6e301b..edba4bf 100644 --- a/packages/training/src/index.ts +++ b/packages/training/src/index.ts @@ -1,12 +1,15 @@ export { configureTraining, + instrumentTrainable, provideTrainingDefaults, trainable, training, + wrapTrainable, } from "./training.js"; export type { CaptureSettings, CandidateEvalConfig, + EvolutionSettings, EvolveInput, EvolveResult, OptimizeInput, diff --git a/packages/training/src/training.ts b/packages/training/src/training.ts index 4224406..6182676 100644 --- a/packages/training/src/training.ts +++ b/packages/training/src/training.ts @@ -58,9 +58,20 @@ export interface TracingSettings { readonly attributes?: Attributes; } +/** Background code evolution driven by captured traffic; disabled unless enabled here + * or via `ts-autocode/register`. Rewrites still pass the full gate before applying. */ +export interface EvolutionSettings { + readonly enabled?: boolean; + readonly minTraces?: number; + readonly objective?: string; + readonly evaluation?: Omit; + readonly onEvolved?: (result: EvolveResult) => void; +} + export interface TrainingSettings { readonly engine?: TrainingEngine; readonly executor?: ImplementationExecutor; + readonly evolution?: EvolutionSettings; readonly source?: SourceSettings; readonly store?: TrainingStore; readonly secrets?: SecretProvider; @@ -70,7 +81,7 @@ export interface TrainingSettings { readonly tracing?: TracingSettings; readonly idFactory?: () => string; readonly now?: () => Date; - readonly onError?: (error: unknown, phase: "capture" | "store") => void; + readonly onError?: (error: unknown, phase: "capture" | "store" | "evolve") => void; } export interface OptimizeInput { @@ -191,6 +202,42 @@ class TrainingRuntime implements Training { return executor; } + readonly #evolutionState = new Map(); + + #maybeEvolve(token: TrainableToken): void { + const evolution = this.#settings.evolution ?? defaultProviders.evolution; + if (evolution?.enabled !== true) return; + const state = this.#evolutionState.get(token.id) ?? { running: false, queued: false, attempted: 0 }; + this.#evolutionState.set(token.id, state); + if (state.running) { + state.queued = true; + return; + } + state.running = true; + void (async () => { + await this.flush(); + const minTraces = Math.max(1, evolution.minTraces ?? 3); + const successes = (await this.#store.list(token.id)).filter((record) => record.succeeded).length; + if (successes < state.attempted + minTraces) return; + state.attempted = successes; + const result = await this.evolve({ + trainable: token, + objective: evolution.objective ?? "Preserve behavior observed in successful runtime traces", + minTraces, + ...(evolution.evaluation === undefined ? {} : { evaluation: evolution.evaluation }), + }); + evolution.onEvolved?.(result); + })() + .catch((error) => this.#settings.onError?.(error, "evolve")) + .finally(() => { + state.running = false; + if (state.queued) { + state.queued = false; + this.#maybeEvolve(token); + } + }); + } + async records(identity?: TrainableIdentity): Promise { await this.flush(); return this.#store.list(identity === undefined ? undefined : toTrainableToken(identity).id); @@ -509,6 +556,7 @@ class TrainingRuntime implements Training { }), }; this.#enqueue(this.#store.append(record)); + if (error === undefined) this.#maybeEvolve(token); } catch (captureError) { this.#settings.onError?.(captureError, "capture"); } @@ -537,6 +585,7 @@ export function configureTraining(settings: TrainingSettings = {}): Training { export interface TrainingProviders { readonly engine?: () => TrainingEngine; readonly executor?: ImplementationExecutor; + readonly evolution?: EvolutionSettings; } let defaultProviders: TrainingProviders = {}; @@ -563,6 +612,8 @@ export const training: Training = Object.freeze({ flush: () => runtime().flush(), }); +const wrappedMarker = Symbol.for("ts-autocode.wrapped"); + /** Decorator form: `@trainable()`. Identity is inferred from the decorated class and * method; pass a symbol (for example `defineTrainable("Router.route").symbol`) only * to override the inferred id. */ @@ -577,11 +628,44 @@ export function trainable(identity?: symbol): TrainableDecorator { ) { const name = String(context.name); let token = explicit; - return function (this: This, ...args: Args): Result { + const wrapped = function (this: This, ...args: Args): Result { token ??= defineTrainable(`${inferredClassName(this) ?? "Anonymous"}.${name}`); return runtime().invoke(this, method, args, token, name); }; + return markWrapped(wrapped); + }; +} + +/** Load-time instrumentation (`ts-autocode/register`): capture-wrap a directive-marked + * function. Idempotent — already-wrapped functions (decorator or register) pass through. */ +export function wrapTrainable unknown>(fn: F, id: string): F { + if ((fn as Partial>)[wrappedMarker]) return fn; + const token = defineTrainable(id); + const name = fn.name || token.id; + const method = fn as unknown as (this: unknown, ...args: unknown[]) => unknown; + const wrapped = function (this: unknown, ...args: unknown[]): unknown { + return runtime().invoke(this, method, args, token, name); }; + Object.defineProperty(wrapped, "name", { value: name, configurable: true }); + return markWrapped(wrapped) as unknown as F; +} + +/** Load-time instrumentation (`ts-autocode/register`): capture-wrap a directive-marked + * class method in place. Idempotent and tolerant of missing members. */ +export function instrumentTrainable( + owner: abstract new (...args: never[]) => unknown, + methodName: string, + id: string, +): void { + const container = (Object.hasOwn(owner, methodName) ? owner : owner.prototype) as Record; + const method = container?.[methodName]; + if (typeof method !== "function") return; + container[methodName] = wrapTrainable(method as (...args: never[]) => unknown, id); +} + +function markWrapped(fn: F): F { + Object.defineProperty(fn, wrappedMarker, { value: true }); + return fn; } function inferredClassName(thisValue: unknown): string | undefined { diff --git a/packages/training/test/training.test.ts b/packages/training/test/training.test.ts index 4945d2e..71f3ee4 100644 --- a/packages/training/test/training.test.ts +++ b/packages/training/test/training.test.ts @@ -8,8 +8,10 @@ import * as publicApi from "../src/index.js"; import { configureTraining, defineTrainable, + instrumentTrainable, trainable, training as defaultTraining, + type EvolveResult, type ImplementationExecutor, type TrainingEngine, } from "../src/index.js"; @@ -60,6 +62,20 @@ describe("trainable identity", () => { it("rejects non-symbol decorator identities", () => { expect(() => trainable("Router.route" as never)).toThrow("must be a symbol"); }); + + it("instruments classes in place for capture without the decorator", async () => { + configureTraining({ tracing: { enabled: false } }); + class Plain { + route(input: string): string { return input; } + } + instrumentTrainable(Plain, "route", "Plain.route"); + instrumentTrainable(Plain, "route", "Plain.route"); + + expect(new Plain().route("billing")).toBe("billing"); + const records = await defaultTraining.records("Plain.route"); + expect(records).toHaveLength(1); + expect(records[0]?.trainableId).toBe("Plain.route"); + }); }); describe("trainable method capture", () => { @@ -176,6 +192,44 @@ describe("training execution", () => { expect(await readFile(artifact, "utf8")).toContain('"use training"'); }); + it("evolves automatically from runtime traffic when evolution is enabled", async () => { + const directory = await mkdtemp(join(tmpdir(), "ts-autocode-auto-")); + const artifact = join(directory, "auto.ts"); + await writeFile(artifact, `export function autoNormalize(input: string): string { + "use training"; + return input; +}\n`); + let resolveEvolved!: (result: EvolveResult) => void; + const evolved = new Promise((resolve) => { resolveEvolved = resolve; }); + const errors: unknown[] = []; + configureTraining({ + engine: { id: "auto-test", optimize: async () => ({ implementation: "return input.toUpperCase();" }) }, + executor: functionExecutor, + source: { files: [artifact] }, + tracing: { enabled: false }, + onError: (error) => errors.push(error), + evolution: { + enabled: true, + minTraces: 2, + evaluation: { outputDir: join(directory, "agentv") }, + onEvolved: (result) => resolveEvolved(result), + }, + }); + class AutoNormalizer { + normalize(input: string): string { return input.toUpperCase(); } + } + instrumentTrainable(AutoNormalizer, "normalize", "autoNormalize"); + const normalizer = new AutoNormalizer(); + normalizer.normalize("alpha"); + normalizer.normalize("beta"); + + const result = await evolved; + expect(errors).toEqual([]); + expect(result.training.outcome).toBe("ready"); + expect(await readFile(artifact, "utf8")).toContain("return input.toUpperCase();"); + expect(await readFile(artifact, "utf8")).toContain('"use training"'); + }); + it("waives the conformance requirement instead of rejecting every candidate", async () => { const directory = await mkdtemp(join(tmpdir(), "ts-autocode-conformance-")); const artifact = join(directory, "echo.ts"); diff --git a/src/register.ts b/src/register.ts new file mode 100644 index 0000000..e09d018 --- /dev/null +++ b/src/register.ts @@ -0,0 +1,30 @@ +import { registerHooks } from "node:module"; +import { fileURLToPath } from "node:url"; + +import { instrumentTrainable, provideTrainingDefaults, wrapTrainable } from "ts-autocode-training"; + +import { augmentSource, instrumentKey } from "./register/hook.js"; +// Importing the package entry wires the Ax engine and executor defaults. +import "./index.js"; + +(globalThis as Record)[Symbol.for(instrumentKey)] = Object.freeze({ + method: instrumentTrainable, + wrap: wrapTrainable, +}); + +const evolveFlag = (process.env["TS_AUTOCODE_EVOLVE"] ?? "").trim().toLowerCase(); +if (!["0", "false", "off"].includes(evolveFlag)) { + provideTrainingDefaults({ evolution: { enabled: true } }); +} + +registerHooks({ + load(url, context, nextLoad) { + const result = nextLoad(url, context); + if (!url.startsWith("file:") || url.includes("/node_modules/")) return result; + const source = result.source; + if (typeof source !== "string" && !(source instanceof Uint8Array)) return result; + const text = typeof source === "string" ? source : Buffer.from(source).toString("utf8"); + const augmented = augmentSource(text, fileURLToPath(url)); + return augmented === text ? result : { ...result, source: augmented }; + }, +}); diff --git a/src/register/hook.ts b/src/register/hook.ts new file mode 100644 index 0000000..2f53069 --- /dev/null +++ b/src/register/hook.ts @@ -0,0 +1,28 @@ +import { discoverInSource } from "ts-autocode-training"; + +export const instrumentKey = "ts-autocode.instrument"; + +/** Appends guarded instrumentation for every `"use training"` function so the + * register runtime can capture calls without any consumer code. Pure: returns + * the source unchanged when there is nothing to instrument or parsing fails. */ +export function augmentSource(source: string, path: string): string { + if (!source.includes("use training")) return source; + let lines: string[]; + try { + lines = discoverInSource(source, path).flatMap((target) => { + if (target.className) { + return `if (typeof ${target.className} === "function") __tsAutocodeInstrument.method(${target.className}, ${JSON.stringify(target.methodName)}, ${JSON.stringify(target.id)});`; + } + return `if (typeof ${target.methodName} === "function") ${target.methodName} = __tsAutocodeInstrument.wrap(${target.methodName}, ${JSON.stringify(target.id)});`; + }); + } catch { + return source; + } + if (lines.length === 0) return source; + return `${source} +;const __tsAutocodeInstrument = globalThis[Symbol.for(${JSON.stringify(instrumentKey)})]; +if (__tsAutocodeInstrument) { +${lines.join("\n")} +} +`; +} diff --git a/test/register.test.ts b/test/register.test.ts new file mode 100644 index 0000000..f53bbd9 --- /dev/null +++ b/test/register.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from "vitest"; + +import { augmentSource } from "../src/register/hook.js"; + +describe("register source augmentation", () => { + it("appends guarded instrumentation for directive-marked classes and functions", () => { + const source = `export class Router { + route(input) { + "use training"; + return input; + } +} +export function normalize(input) { + "use training"; + return input; +} +`; + const augmented = augmentSource(source, "/app/router.js"); + + expect(augmented.startsWith(source)).toBe(true); + expect(augmented).toContain('__tsAutocodeInstrument.method(Router, "route", "Router.route");'); + expect(augmented).toContain('normalize = __tsAutocodeInstrument.wrap(normalize, "normalize");'); + expect(augmented).toContain('globalThis[Symbol.for("ts-autocode.instrument")]'); + }); + + it("returns sources without trainables unchanged, including on parse-adjacent content", () => { + expect(augmentSource("export const x = 1;\n", "/app/x.js")).toBe("export const x = 1;\n"); + const mention = "// mentions use training in a comment only\nexport const y = 2;\n"; + expect(augmentSource(mention, "/app/y.js")).toBe(mention); + }); +}); From 37a693fa6da86222fdddb88d17ae580f7194bc2f Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 12 Jul 2026 01:50:41 +0000 Subject: [PATCH 6/8] chore: add Ax library skill docs for coding agents Fifteen @ax-llm/ax SKILL.md references (providers, signatures, generators, agents, flows, GEPA) so agent tooling working in this repo has current Ax v23 API documentation. Content reviewed: documentation and code samples only. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X2LxfjsASWmdVoFANEURWR --- .claude/skills/ax-agent-context/SKILL.md | 43 ++ .../skills/ax-agent-memory-skills/SKILL.md | 438 ++++++++++++ .../skills/ax-agent-observability/SKILL.md | 353 ++++++++++ .claude/skills/ax-agent-optimize/SKILL.md | 360 ++++++++++ .claude/skills/ax-agent-rlm/SKILL.md | 501 ++++++++++++++ .claude/skills/ax-agent/SKILL.md | 631 ++++++++++++++++++ .claude/skills/ax-ai/SKILL.md | 382 +++++++++++ .claude/skills/ax-audio/SKILL.md | 373 +++++++++++ .claude/skills/ax-flow/SKILL.md | 442 ++++++++++++ .claude/skills/ax-gen/SKILL.md | 495 ++++++++++++++ .claude/skills/ax-gepa/SKILL.md | 264 ++++++++ .claude/skills/ax-llm/SKILL.md | 339 ++++++++++ .claude/skills/ax-playbook/SKILL.md | 95 +++ .claude/skills/ax-refine/SKILL.md | 81 +++ .claude/skills/ax-signature/SKILL.md | 306 +++++++++ 15 files changed, 5103 insertions(+) create mode 100644 .claude/skills/ax-agent-context/SKILL.md create mode 100644 .claude/skills/ax-agent-memory-skills/SKILL.md create mode 100644 .claude/skills/ax-agent-observability/SKILL.md create mode 100644 .claude/skills/ax-agent-optimize/SKILL.md create mode 100644 .claude/skills/ax-agent-rlm/SKILL.md create mode 100644 .claude/skills/ax-agent/SKILL.md create mode 100644 .claude/skills/ax-ai/SKILL.md create mode 100644 .claude/skills/ax-audio/SKILL.md create mode 100644 .claude/skills/ax-flow/SKILL.md create mode 100644 .claude/skills/ax-gen/SKILL.md create mode 100644 .claude/skills/ax-gepa/SKILL.md create mode 100644 .claude/skills/ax-llm/SKILL.md create mode 100644 .claude/skills/ax-playbook/SKILL.md create mode 100644 .claude/skills/ax-refine/SKILL.md create mode 100644 .claude/skills/ax-signature/SKILL.md diff --git a/.claude/skills/ax-agent-context/SKILL.md b/.claude/skills/ax-agent-context/SKILL.md new file mode 100644 index 0000000..738baea --- /dev/null +++ b/.claude/skills/ax-agent-context/SKILL.md @@ -0,0 +1,43 @@ +--- +name: ax-agent-context +description: This skill helps an LLM pick the right AxAgent context tool for a job - contextMap for recurring corpora, contextPolicy presets for within-run trajectory compaction, agent.optimize for offline GEPA instruction/demo tuning, agent.playbook for an evolving context playbook (offline evolve + online update), and recall/memories + skills for per-turn retrieval. Use when the user asks "which context feature should I use", confuses contextMap with contextPolicy or memory, or wants a decision guide for long-context agents. For contextPolicy/contextMap codegen use ax-agent-rlm; for recall/skills use ax-agent-memory-skills; for agent.optimize or agent.playbook use ax-agent-optimize. +version: "23.0.0" +--- + +# AxAgent Context Selection (@ax-llm/ax) + +Use this skill to route a context-management need to the right AxAgent tool, then open the matching codegen skill. AxAgent manages four distinct context objects; choosing the wrong one is the usual mistake. Do not write tutorial prose; pick the tool and hand off. + +## Pick The Right Context Tool + +| Need | Object | Scope | Use | Next skill | +| --- | --- | --- | --- | --- | +| Many tasks over the same large corpus (repo, doc set, dataset) | Context map | recurring corpus, persists across runs | `contextMap` | `ax-agent-rlm` | +| One long run whose own history must stay under control | Trajectory compaction | this run only | `contextPolicy: { preset, budget }` | `ax-agent-rlm` | +| Evolve task strategy from examples or live feedback | Context playbook | a stage, offline + online | `agent.playbook(...)` | `ax-agent-optimize` | +| Tune the prompt/instructions/demos offline | Instruction text | a program, offline | `agent.optimize(...)` (GEPA) | `ax-agent-optimize` | +| Pull task-relevant facts or guides for a turn | Retrieval | one turn | `recall(...)` / skills | `ax-agent-memory-skills` | + +## Defaults + +- Recurring corpus + many different questions -> `contextMap` (persistent orientation cache). +- One long multi-turn run with prompt pressure -> `contextPolicy: { preset: 'checkpointed', budget: 'balanced' }`; move to `lean` for very long runs with strong models, `full` for short tasks or weak models. +- Evolve a context playbook -> `agent.playbook(...)` (offline from examples, or online from live feedback). +- Tune instructions/demos offline -> `agent.optimize(...)` (GEPA). +- Fetch facts or guides on demand -> `recall(...)` for memories, `discover({ skills })` for skill guides. +- A single oversized input value (a pasted doc, a big JSON blob) -> do nothing; `autoUpgrade` (ON by default) keeps it runtime-only with a prompt preview. Reach for `contextFields` only when you want a specific inline policy or the value is a large required non-string field. See `ax-agent-rlm`. + +## Anti-Patterns + +- Do not use `contextMap` to compress a single run's history. That is `contextPolicy`. +- Do not use `contextPolicy` to carry knowledge across runs. That is `contextMap`. +- Do not hand-build a strategy playbook in the prompt. Evolve it with `agent.playbook(...)`. +- Do not stuff a whole corpus into the prompt every run. Use a context map plus on-demand `recall(...)`. +- Do not confuse runtime skills (`discover({ skills })` guides) with these installable codegen skills. + +## See Also + +- `ax-agent-rlm` - contextPolicy presets, context maps, and runtime sessions. +- `ax-agent-memory-skills` - recall, memories, and dynamic skill loading. +- `ax-agent-optimize` - GEPA via `agent.optimize(...)` and the context playbook via `agent.playbook(...)`. +- `ax-agent` - core agent shape and the final/clarification protocol. diff --git a/.claude/skills/ax-agent-memory-skills/SKILL.md b/.claude/skills/ax-agent-memory-skills/SKILL.md new file mode 100644 index 0000000..59de391 --- /dev/null +++ b/.claude/skills/ax-agent-memory-skills/SKILL.md @@ -0,0 +1,438 @@ +--- +name: ax-agent-memory-skills +description: This skill helps an LLM generate correct AxAgent memory retrieval, context-map, and dynamic skill-loading code using @ax-llm/ax. Use when the user asks about contextMap, AxAgentContextMap, onMemoriesSearch, memoriesCatalog, recall(...), inputs.memories, onLoadedMemories, onUsedMemories, onSkillsSearch, skillsCatalog, AxAgentCatalogSkill, discover({ skills }), onLoadedSkills, onUsedSkills, preloaded skills, preloading memories at forward time, relevanceRanking hints, loaded memory/skill IDs, or carrying memories across forward() calls. +version: "23.0.0" +--- + +# AxAgent Memory And Skills Rules (@ax-llm/ax) + +Use this skill when an agent needs a persistent context map, task-relevant memory retrieval, or skill guides loaded into the executor prompt on demand. For ordinary agent setup use `ax-agent`. For RLM runtime policy use `ax-agent-rlm`. For callbacks and telemetry use `ax-agent-observability`. + +## Use These Defaults + +- Use a static `skillsCatalog` / `memoriesCatalog` when the skill guides or memories fit in a plain array — Ax then backs `discover({ skills })` / `recall(...)` with a built-in deterministic local search and no host search code is needed. +- Use `onSkillsSearch` / `onMemoriesSearch` when retrieval needs a real backend (vector DB, BM25 service, KV). A host callback always takes precedence over the catalog's built-in search. +- Use `contextMap` when repeated runs inspect the same long external context and should accumulate a small orientation cache automatically. +- `recall(...)` is available to distiller and executor stages when `onMemoriesSearch` or a non-empty `memoriesCatalog` is set. +- `discover({ skills })` is available to the executor when `onSkillsSearch` or a non-empty `skillsCatalog` is set. +- With `skillsCatalog`, the executor prompt also gains a static `### Available Skills` index (id + name + description), so skill discovery is targeted instead of blind. +- Both `recall(...)` and `discover({ skills })` return `void`. The loaded content appears on the next turn. +- Use `onLoadedMemories` / `onLoadedSkills` to observe what got loaded. +- Use `onUsedMemories` / `onUsedSkills` to track what the actor says it actually relied on. +- Child agents do not inherit memory or skills search callbacks; wire them explicitly on every agent that needs the capability. + +## Context Map + +Use `contextMap` when repeated runs ask different questions over the same long context, document set, or repository. The map is prompt-resident orientation knowledge: structure, concepts, constants, parsing schema, reusable aggregate results, and concrete error patterns. It is not a task-specific answer cache. + + Runnable example: [`src/examples/rlm-context-map-live.ts`](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/rlm-context-map-live.ts) demonstrates a provider-backed context-map update, `onUpdate` snapshot persistence, finite evolve, and frozen map reuse. + +When `contextMap` is configured: + +- Ax injects the current map into the distiller prompt. +- Ax updates the map once after each successful completed `forward(...)`. +- By default the map evolves forever. For a finite warmup, create the map with `{ infiniteEvolve: false, evolveSteps: N }`; after `N` successful updates it is still injected but no longer updated. +- Failed runs, aborts, and clarification requests do not update the map. +- Use `onUpdate` to persist `result.map.snapshot()` outside the agent. + +```typescript +import { agent, AxAgentContextMap } from '@ax-llm/ax'; + +const map = new AxAgentContextMap(savedSnapshot, { + maxChars: 4000, + infiniteEvolve: false, + evolveSteps: 10, +}); + +const myAgent = agent('context:string, query:string -> answer:string', { + contextFields: ['context'], + contextMap: { + map, + onUpdate: ({ map }) => saveSnapshot(map.snapshot()), + }, +}); +``` + +Types: + +```typescript +type AxAgentContextMapConfig = { + map?: AxAgentContextMap | AxAgentContextMapSnapshot | string; + onUpdate?: (result: AxAgentContextMapUpdateResult) => void | Promise; +}; + +type AxAgentContextMapOptions = { + maxChars?: number; + infiniteEvolve?: boolean; + evolveSteps?: number; +}; +``` + +## Memory Search + +Use `onMemoriesSearch` when the agent needs to pull task-relevant context such as user preferences, prior decisions, project facts, or past conversations from an external store (vector DB, BM25, KV). The actor decides what to load, when, and how much. + +When `onMemoriesSearch` is set, the distiller and executor stages gain: + +1. An `inputs.memories` field. In JS this is an array of `{ id, content }` entries the actor reads directly. In the prompt, the same entries render as markdown blocks with `ID: \`...\`` lines, matching the Loaded Skills ID style. Each `content` is opaque markdown; frontmatter is not parsed. +2. A `recall(searches: string[]): void` global the actor `await`s to load more entries. Recalled entries are appended to `inputs.memories` and visible from the next turn onward. `recall()` returns nothing. + +The responder stage does not receive memories. + +### Enabling + +```typescript +import { agent } from '@ax-llm/ax'; +import type { AxAgentMemoriesSearchFn } from '@ax-llm/ax'; + +const onMemoriesSearch: AxAgentMemoriesSearchFn = async ( + searches, + alreadyLoaded +) => { + // `searches` is the full array passed to recall(...). Batch your + // store lookup in one round-trip. + // `alreadyLoaded` is the current inputs.memories snapshot. Filter + // against it to skip duplicates. + const skip = new Set(alreadyLoaded.map((m) => m.id)); + const fresh = await myVectorDB.searchBatch(searches, { topK: 3 }); + return fresh.filter((m) => !skip.has(m.id)); +}; + +const myAgent = agent('task:string -> answer:string', { + contextFields: [], + onMemoriesSearch, +}); +``` + +Each memory result must be: + +```typescript +type AxAgentMemoryResult = { + id: string; + content: string; +}; +``` + +### Static catalog (no callback) + +If the memory set fits in a plain array, skip the callback entirely: pass `memoriesCatalog` and Ax backs `recall(...)` with a built-in deterministic local search (idf-weighted token overlap over `id` + content; not regex, not embeddings). The `alreadyLoaded` contract is preserved — entries already on `inputs.memories` are excluded before ranking. + +```typescript +const myAgent = agent('task:string -> answer:string', { + contextFields: [], + memoriesCatalog: [ + { id: 'deploy-window', content: 'Prod deploys only on Tuesday afternoons.' }, + { id: 'user-prefs', content: 'User prefers concise answers.' }, + ], +}); +``` + +Rules: + +- If both `memoriesCatalog` and `onMemoriesSearch` are set, the host callback handles all `recall(...)` searches; the catalog still powers the advisory `relevanceRanking` hint. +- Catalog content is NOT preloaded into the prompt; entries load only when recalled. +- The built-in search is lexical. For semantic retrieval over large stores, supply `onMemoriesSearch` instead. + +### Preloading memories at forward time + +To seed specific memories for one run (no recall round-trip), pass them as the `memories` input value. They render on `inputs.memories` from the first turn and merge with anything recalled later (deduped by `id`). + +```typescript +await myAgent.forward(ai, { + task: 'Plan the deploy', + memories: [{ id: 'deploy-window', content: 'Prod deploys only on Tuesday afternoons.' }], +}); +``` + +### Actor usage + +```javascript +// Turn 1: kick off one batched lookup. +await recall(['user preferences', 'project constraints']); + +// Turn 2+: matched entries are now visible on inputs.memories. +const prefs = inputs.memories.find((m) => m.id === 'user-prefs-v2'); +``` + +Rules: + +- Pass all memory queries in one `await recall([...])` call. +- Do not loop `recall()` calls or wrap them in `Promise.all(...)`. +- Read `inputs.memories` on the next turn to see what landed. +- `recall()` invokes `onMemoriesSearch` with `(searches, alreadyLoaded)` and returns `void`. +- Results land on `inputs.memories` for subsequent turns and render in the prompt as: + +```markdown +### Memory + +ID: `mem:user-prefs-v2` + +... +``` + +- Entries are deduped by `id` (last-write-wins) and sorted by `id` for prefix-cache stability. +- Memories loaded by the distiller thread automatically to the executor. No second `recall()` is needed for those entries. +- `recall()` may be called multiple times across turns; results accumulate for that run. +- `inputs.memories` lifetime is one `.forward()` call. It resets between calls. + +## Carrying Memories Across `.forward()` Calls + +To preserve continuity across calls, persist memories in your store and recall them again on the next call. If you want to replay anything loaded on a prior run, observe loads with `onLoadedMemories`. + +```typescript +const carried = new Map(); + +const myAgent = agent('task:string -> answer:string', { + contextFields: [], + onMemoriesSearch: async (searches) => { + const fresh = await myVectorDB.searchBatch(searches, { topK: 3 }); + const carriedAsResults = [...carried.entries()].map(([id, content]) => ({ + id, + content, + })); + return [...carriedAsResults, ...fresh]; + }, + onLoadedMemories: (results) => { + for (const r of results) carried.set(r.id, r.content); + }, +}); +``` + +## Skills Search + +Use `onSkillsSearch` when the agent needs to load skill guides such as usage instructions, runbooks, or domain conventions into the executor's system prompt on demand. The actor decides which skills to fetch and when, so you do not pre-render every skill into every prompt. + +When `onSkillsSearch` is set, the distiller and executor stages gain: + +1. A "Loaded Skills" section in the system prompt that renders matched skill bodies with stable `ID:` values sorted by `id`. +2. A `discover({ skills })` path the actor `await`s to load more skills. Loaded entries appear in the next turn's prompt. `discover(...)` returns nothing. + +Skills the distiller loads carry over to the executor automatically. The responder does not see skills. + +### Enabling + +```typescript +import { agent } from '@ax-llm/ax'; +import type { AxAgentSkillsSearchFn } from '@ax-llm/ax'; + +// Each result is { id?: string; name: string; content: string }. +// If id is omitted, Ax falls back to name. +const onSkillsSearch: AxAgentSkillsSearchFn = async (searches) => { + return mySkillStore.resolveBatch(searches, { + // Recommended backend order: exact id, exact name, then broader search. + // This lets the actor pass one simple string and keeps lookup policy host-side. + strategy: ['id', 'name', 'search'], + topK: 2, + }); +}; + +const myAgent = agent('task:string -> answer:string', { + contextFields: [], + onSkillsSearch, +}); +``` + +Each skill result is: + +```typescript +type AxAgentSkillResult = { + id?: string; + name: string; + content: string; +}; +``` + +### Static catalog (no callback) + +If the skill set fits in a plain array, skip the callback entirely: pass `skillsCatalog` and Ax backs `discover({ skills })` with a built-in deterministic local search (idf-weighted token overlap over `id` + `name`×2 + `description`×2 + the first 600 chars of `content`; not regex, not embeddings). The executor prompt also gains a static, cache-stable `### Available Skills` index (id + name + description, sorted by id), so the actor searches by known ids instead of guessing. + +```typescript +import type { AxAgentCatalogSkill } from '@ax-llm/ax'; + +const catalog: AxAgentCatalogSkill[] = [ + { + id: 'release-checklist', + name: 'Release checklist', + description: 'Steps for shipping a package release safely', // high-signal for matching + content: '1. Bump version\n2. Run tests\n3. Tag and publish', + }, +]; + +const myAgent = agent('task:string -> answer:string', { + contextFields: [], + skillsCatalog: catalog, +}); +``` + +```typescript +type AxAgentCatalogSkill = { + id: string; + name: string; + description?: string; + content: string; +}; +``` + +Rules: + +- If both `skillsCatalog` and `onSkillsSearch` are set, the host callback handles all `discover({ skills })` searches; the catalog still powers the `### Available Skills` index and the advisory `relevanceRanking` hint. +- Catalog content is NOT preloaded into the prompt (unlike `skills`); entries load only when matched. Use `skills` for guides that must always be in context, `skillsCatalog` for a larger set loaded on demand. +- The built-in search is lexical. For semantic retrieval over large stores, supply `onSkillsSearch` instead. + +### Actor usage + +```javascript +// Pass all skill queries in one call. +await discover({ skills: ['release-checklist', 'incident-response'] }); + +// Next turn: loaded skill bodies render under the "Loaded Skills" +// system-prompt section. +``` + +Rules: + +- `discover({ skills })` invokes `onSkillsSearch` with the raw search strings and returns `void`. +- Resolve each raw string backend-side: prefer an exact `id` match, then an exact `name` match, then fuzzy/full-text search. The actor should not have to choose `id:` vs `name:` syntax. +- Matched skills land under "Loaded Skills" for the next turn. +- Entries are deduped by `id` (last-write-wins) and sorted by `id` for prefix-cache stability. +- If a skill result omits `id`, its trimmed `name` is used as the id for backwards compatibility. +- Skills persist on the agent's `currentSkillsPromptState` across `.forward()` calls, unlike memories. +- Use `agent.getState()` / `setState(...)` to serialize/restore loaded skills. +- `discover({ skills })` may be called multiple times across turns. Within one turn, batch all skill queries in a single call. +- Child agents do not inherit `onSkillsSearch`; wire it explicitly per agent. + +## Preloading Skills + +If the caller already knows which skills are relevant, pass them up front instead of round-tripping through `discover({ skills })`. + +- Init-time: `skills` on `AxAgentOptions` seeds the executor prompt at agent creation. They survive `setState(...)` resets. +- Forward-time: `skills` on `forward(ai, values, { skills })` merge in at the start of that call. Distiller and responder ignore forward-time skills. + +Both accept the same shape `onSkillsSearch` returns: `readonly AxAgentSkillResult[]`. Forward-time skills override init-time skills by `id`. `onLoadedSkills` is not fired for preset skills; that callback is for runtime `discover({ skills })` analytics. + +```typescript +const releaseAgent = agent('task:string -> answer:string', { + contextFields: [], + skills: [ + { + id: 'release-checklist', + name: 'release-checklist', + content: '...', + }, + ], +}); + +await releaseAgent.forward( + ai, + { task: 'Prepare release notes' }, + { + skills: [ + { + id: 'incident-response', + name: 'incident-response', + content: '...', + }, + ], + } +); +``` + +You can use `skills` without setting `onSkillsSearch` at all. That is useful for static guides where the actor never needs to fetch more. + +## Advisory Relevance Hints (`relevanceRanking`) + +`relevanceRanking` is ON by default — leave it unset; set `relevanceRanking: false` to opt out. The default was flipped after its A/B gate passed (substance-judged, 49 runs per variant per model: small-model first-lookup precision 24%→90% and answer accuracy 14%→29%; frontier-model control accuracy 63%→88% with fewer turns). The generated language ports implement the same advisory hint contract through AxIR Core. + +When enabled, a deterministic local ranker scores the agent's discoverable capabilities against the task once per `forward(...)` and injects a short advisory `### Likely Relevant` shortlist into the executor turn — modules (needs `functionDiscovery`), catalog skills (needs `skillsCatalog`), and catalog memories (needs `memoriesCatalog`). The hint is non-authoritative: the full lists still apply and the actor may `discover`/`recall` anything else. + +```typescript +const myAgent = agent('task:string -> answer:string', { + contextFields: [], + functionDiscovery: true, + skillsCatalog: catalog, + relevanceRanking: true, // or { topK: 3, minScore: 0.08 } +}); +``` + +Rules: + +- Default is ON across TypeScript and generated language ports; domains still self-gate on their prerequisites (`functionDiscovery` for modules, catalogs for skills/memories), so agents without those see no change. Everything else in this skill (catalog search, the Available Skills index) is independent of the flag. +- The shortlist rides a dynamic, non-cached prompt field; the cached system prompt stays byte-stable across tasks. +- On low confidence the ranker emits nothing rather than guessing. +- Memory hint entries include an ~80-char content snippet; very short memories may be usable from the hint alone without a `recall(...)` (such use is not visible to `onUsedMemories`). +- Observe outcomes via the `relevance_ranking` context event (see `ax-agent-observability`). + +## Loaded And Used Tracking + +`onLoadedMemories` reports what `recall(...)` loaded. `onLoadedSkills` reports what `discover({ skills })` loaded. To track what the actor says it actually relied on, use `onUsedMemories` / `onUsedSkills`. + +```typescript +const used: AxAgentUsedMemory[] = []; + +await myAgent.forward( + ai, + { task: 'Make a personal plan' }, + { + onUsedMemories: (items) => used.push(...items), + } +); + +used; // [{ id, reason, stage }] +``` + +Rules: + +- The actor can only report memory IDs already present in `inputs.memories`. +- The actor can only report skill IDs already present in Loaded Skills. +- Unknown values are dropped. +- When tracking is enabled, the actor sees `await used(id, reason?)`; this is the actor-side declaration mechanism. +- `used(...)` resolves against loaded memory IDs and loaded skill IDs. +- If memory IDs and skill IDs can collide, namespace them in your application, for example `mem:abc` and `skill:planning`. + +Types: + +```typescript +onMemoriesSearch?: AxAgentMemoriesSearchFn; +onLoadedMemories?: ( + results: readonly AxAgentMemoryResult[] +) => void | Promise; +onUsedMemories?: ( + usedMemories: readonly AxAgentUsedMemory[] +) => void | Promise; + +onSkillsSearch?: AxAgentSkillsSearchFn; +onLoadedSkills?: ( + results: readonly AxAgentSkillResult[] +) => void | Promise; +onUsedSkills?: ( + usedSkills: readonly AxAgentUsedSkill[] +) => void | Promise; + +contextMap?: AxAgentContextMapConfig; +skills?: readonly AxAgentSkillResult[]; +skillsCatalog?: readonly AxAgentCatalogSkill[]; +memoriesCatalog?: readonly AxAgentMemoryResult[]; +relevanceRanking?: boolean | { topK?: number; minScore?: number }; +``` + +## Examples + +Fetch this for full working code: + +- [RLM Memories and Skills](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/rlm-memories-and-skills.ts) - `onMemoriesSearch` + `recall()` and `onSkillsSearch` + `discover({ skills })` with load observability and actual usage tracking via `onUsedMemories` / `onUsedSkills` +- [Skills + Memory Ops Assistant](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/typescript/long-agents/skills-and-memory-assistant.ts) - an on-call assistant that recalls past decisions from a memory store and loads the right runbook skill on demand (also ported to Python, Go, Rust, Java, and C++ under `src/examples//long-agents/`). All six languages support the native `onMemoriesSearch` / `onSkillsSearch` host callbacks, passed in the agent options at construction (Go/Java use native function values, Rust a `agent_with_search_callbacks` constructor, C++ a `register_*_search` helper); a static `memory_search_results` / `skill_search_results` config is also available. + +## Do Not Generate + +- Do not assign the result of `await recall(...)` or `await discover(...)`; both return `void`. +- Do not call `recall()` from the responder stage. +- Do not call `discover({ skills })` from the responder stage. +- Do not loop `recall()` calls or wrap them in `Promise.all(...)`. +- Do not loop `discover()` calls or wrap them in `Promise.all(...)`. +- Do not assume child agents inherit `onMemoriesSearch` or `onSkillsSearch`. +- Do not pass `onMemoriesSearch` results via shared fields as a workaround; use `recall(...)`. +- Do not assume `inputs.memories` persists across `.forward()` calls. +- Do not use `onLoadedMemories` / `onLoadedSkills` as proof that the actor relied on an item; use `onUsedMemories` / `onUsedSkills` for actual-use tracking. +- Do not write an `onSkillsSearch` / `onMemoriesSearch` callback that just scans a static array; pass the array as `skillsCatalog` / `memoriesCatalog` instead. +- Do not rely on the built-in catalog search for semantic matching over large stores; it is lexical token overlap — supply a host callback for embeddings/vector search. +- Do not confuse `skills` (always preloaded into the prompt) with `skillsCatalog` (searchable, loaded on demand). diff --git a/.claude/skills/ax-agent-observability/SKILL.md b/.claude/skills/ax-agent-observability/SKILL.md new file mode 100644 index 0000000..23ec76b --- /dev/null +++ b/.claude/skills/ax-agent-observability/SKILL.md @@ -0,0 +1,353 @@ +--- +name: ax-agent-observability +description: This skill helps an LLM generate correct AxAgent observability code using @ax-llm/ax. Use when the user asks about actorTurnCallback, onContextEvent, agentStatusCallback, onFunctionCall, reportSuccess, reportFailure, getChatLog(), getUsage(), resetUsage(), debug traces, progress updates, or telemetry for AxAgent runs. +version: "23.0.0" +--- + +# AxAgent Observability Rules (@ax-llm/ax) + +Use this skill when an agent needs runtime visibility, progress reporting, tracing, usage accounting, or chat-log access. For ordinary agent setup use `ax-agent`. For RLM runtime policy use `ax-agent-rlm`. For memories and dynamic skill loading use `ax-agent-memory-skills`. + +## Choose The Smallest Hook + +- Need a quick prompt/runtime trace during development -> start with `debug: true`. +- Need structured per-turn code, raw runtime result, formatted output, provider thoughts, or actor stage -> use `actorTurnCallback`. +- Need context-pressure and compaction telemetry -> use `onContextEvent`. +- Need real-time task progress emitted by actor code -> use `agentStatusCallback`. +- Need every runtime function call before execution -> use `onFunctionCall`. +- Need model prompts/responses after a run -> use `getChatLog()`. +- Need token usage by actor/responder -> use `getUsage()` and `resetUsage()`. +- Need usage split by context and task stages -> use `getStagedUsage()`. +- Need Ax program traces -> use `getTraces()`. +- Do not add multiple hooks unless the user clearly needs each output stream. + +## Global Runtime Defaults + +OpenTelemetry and debug defaults come from the shared Ax runtime surface: + +```typescript +import { axGlobals, axCreateDefaultColorLogger } from '@ax-llm/ax'; +import { trace } from '@opentelemetry/api'; + +axGlobals.tracer = trace.getTracer('agent-app'); +axGlobals.debug = true; +axGlobals.logger = axCreateDefaultColorLogger(); +``` + +These globals are live defaults for future AI, AxGen, AxFlow, and agent-internal model calls. Per-call or explicitly configured options still override `axGlobals`. Use AxAgent callbacks below when the caller needs structured agent-turn events rather than OpenTelemetry spans or debug logs. + +## Actor Turn Callback + +Use `actorTurnCallback` when the caller needs structured telemetry for each actor turn. + +What it gives you: + +- `code`: the normalized JavaScript code the actor produced +- `stage`: which actor produced the turn (`distiller` or `executor`) +- `result`: the raw untruncated runtime return value from executing that code +- `output`: the formatted action-log output string after Ax normalizes and truncates it for prompt replay +- `thought`: the actor model's `thought` field when `showThoughts` is enabled and the provider returns one +- `executorResult`: the full actor payload returned by the current actor stage, kept under this historical field name for compatibility +- `isError`: whether the execution path for that turn was treated as an error +- `usage`: token usage for this actor turn only +- `model`: model used for this turn when explicitly set through `executorModelPolicy` +- `chatLogMessages`: raw ChatML conversation for this turn, populated only when an actor turn callback is set + +Use it for: + +- debug UIs that want to show code plus raw runtime results +- tracing and analytics +- capturing `thought` for internal diagnostics when supported by the provider +- storing per-turn execution artifacts without scraping the prompt/action log + +Important: + +- `output` is not raw stdout; it is the formatted replay string used in the action log. +- `result` is the raw runtime result before Ax applies type-aware serialization and budget-proportional truncation. +- `thought` is optional and only appears when the underlying `AxGen` call had `showThoughts` enabled and the provider actually returned a thought field. +- `actionLogEntryCount` and `guidanceLogEntryCount` reflect the live log sizes after the turn is processed, including resumed runs. +- `actorTurnCallback` fires for the configured agent instance. Child agents passed through `functions: [...]` should define their own callback if you need their internal actor turns; use `onFunctionCall` on the parent to observe the parent-side child-agent invocation. + +Good pattern: + +```typescript +const supportAgent = agent('query:string -> answer:string', { + contextFields: ['query'], + runtime, + actorTurnCallback: ({ + stage, + turn, + actionLogEntryCount, + guidanceLogEntryCount, + code, + result, + output, + thought, + isError, + usage, + model, + }) => { + console.log({ + turn, + stage, + model, + actionLogEntryCount, + guidanceLogEntryCount, + isError, + code, + rawResult: result, + replayOutput: output, + thought, + usage, + }); + }, + executorOptions: { + model: 'gpt-5.4-mini', + showThoughts: true, + }, +}); +``` + +Callback type: + +```typescript +actorTurnCallback?: (turn: { + stage: 'distiller' | 'executor'; + turn: number; + actionLogEntryCount: number; + guidanceLogEntryCount: number; + executorResult: Record; + code: string; + result: unknown; + output: string; + isError: boolean; + thought?: string; + usage?: AxProgramUsage[]; + model?: string; + chatLogMessages?: ReadonlyArray<{ role: string; content: string }>; +}) => void | Promise; + +actorTurnCallback?: (turn: { + stage: 'distiller' | 'executor'; + turn: number; + actionLogEntryCount: number; + guidanceLogEntryCount: number; + executorResult: Record; + code: string; + result: unknown; + output: string; + isError: boolean; + thought?: string; + usage?: AxProgramUsage[]; + model?: string; + chatLogMessages?: ReadonlyArray<{ role: string; content: string }>; +}) => void | Promise; // deprecated alias +``` + +## Context Event Observability + +Use `onContextEvent` when the caller needs structured telemetry about prompt pressure and compaction. It does not change model behavior directly; it is for logs, evals, and dashboards. + +Events: + +- `budget_check`: character-based prompt pressure before an actor turn, with detailed metrics kept out of the actor prompt +- `checkpoint_created` / `checkpoint_cleared`: checkpoint lifecycle events with covered turns and reason +- `tombstone_created`: compact resolved-error summary creation +- `relevance_ranking`: emitted once per ranked domain per forward when `relevanceRanking` is enabled; carries `domain` (`'modules' | 'skills' | 'memories'`), the `shortlist` (`{ id, score }[]`, most relevant first), and `suppressed` (true when the low-confidence guard emitted no hint) +- `field_auto_promoted`: emitted once per field per run when `autoUpgrade` keeps an oversized undeclared input value runtime-only; carries `fieldName`, `originalChars`, and `promptPreviewChars` (undefined when no inline preview was kept) + +To measure whether the advisory hint helps, join per forward: `relevance_ranking.shortlist` ids against what the actor then loaded — for modules the internal `discover` calls (`onFunctionCall` with `kind: 'internal'`, `name: 'discover'`, `args.request`) plus the module part of external `qualifiedName`s; for skills `onLoadedSkills` / `used(id)`; for memories `onLoadedMemories` / `used(id)`. + +Rules: + +- `contextPressure` in the actor prompt is intentionally compact (`ok`, `watch`, `critical` plus one short instruction). +- Budget metrics are character-based for provider neutrality and are exposed through `onContextEvent`, not the actor prompt. +- Callback errors are swallowed so telemetry cannot break the agent run. +- Do not scrape actor prompts for pressure metrics. + +```typescript +const supportAgent = agent('query:string -> answer:string', { + contextFields: ['query'], + runtime, + contextPolicy: { preset: 'checkpointed', budget: 'balanced' }, + onContextEvent: (event) => { + if (event.kind === 'budget_check') { + console.log(event.pressure, event.mutablePromptChars); + } + }, +}); +``` + +Type: + +```typescript +onContextEvent?: (event: AxAgentContextEvent) => void | Promise; +``` + +## Agent Status Callback + +Use `agentStatusCallback` when the caller wants real-time progress updates from the actor. When set, the actor can call `await reportSuccess(message)` and `await reportFailure(message)` in its JavaScript turns. + +```typescript +const supportAgent = agent('query:string -> answer:string', { + contextFields: ['query'], + runtime, + agentStatusCallback: (message, status) => { + console.log(`[${status}] ${message}`); + }, +}); +``` + +Rules: + +- `agentStatusCallback` receives `(message: string, status: 'success' | 'failed')`. +- When set, the actor prompt automatically includes `reportSuccess(message)` and `reportFailure(message)` as available runtime functions. +- The actor is instructed to keep the user updated on task progress. +- `reportSuccess` and `reportFailure` are reserved runtime names when the callback is configured. +- Child agents inherit the callback via the RLM config. + +Type: + +```typescript +agentStatusCallback?: ( + message: string, + status: 'success' | 'failed' +) => void | Promise; +``` + +## On Function Call + +Use `onFunctionCall` when the caller wants to observe every function call the actor makes from the JS runtime. It fires before the underlying function runs. + +```typescript +const supportAgent = agent('query:string -> answer:string', { + contextFields: ['query'], + runtime, + functions: [helperAgent, lookupOrderTool], + onFunctionCall: ({ name, qualifiedName, args, kind }) => { + console.log(`[${kind}] ${qualifiedName}`, args); + }, +}); +``` + +Rules: + +- Receives `{ name, qualifiedName, args, kind }`. +- `name` is the bare function name, e.g. `'lookupOrder'`. +- `qualifiedName` is the namespaced name as the actor sees it, e.g. `'tools.lookupOrder'`; for un-namespaced runtime globals it equals `name`. +- `args` is the resolved positional/named arguments object (`Record`). +- `kind` is `'external'` for caller-registered `functions`. +- `kind` is `'internal'` for agent-injected globals: child agents, `discover`, `recall`, and `used`. +- Fires once per call, before the function executes. +- Errors thrown inside the callback are swallowed so they cannot break the actor loop. +- This is independent from the DSP-layer `onFunctionCall` on `AxProgramForwardOptions`; that hook is for LLM tool-calls and never fires under AxAgent because AxAgent injects functions as runtime globals. + +Type: + +```typescript +onFunctionCall?: (call: { + name: string; + qualifiedName: string; + args: Record; + kind: 'internal' | 'external'; +}) => void | Promise; +``` + +## Chat Log, Usage, And Traces + +`AxAgent` exposes actor and responder sub-programs. `getChatLog()` returns the same flat `AxChatLogEntry[]` shape as `AxGen` and `AxFlow`; use each entry's optional `name` field to distinguish `distiller`, `executor`, and `responder`. `getUsage()` returns token usage split by actor/responder. + +### getChatLog() + +Returns the full normalized chat history after any `.forward()` call. Each entry is one `ai.chat()` round-trip. Actor stages accumulate one entry per turn; the responder typically has one entry. + +```typescript +const log = myAgent.getChatLog(); + +for (const entry of log) { + console.log(entry.name, entry.model); + for (const msg of entry.messages) { + console.log(`[${msg.role}]`, msg.content); + } +} +``` + +Each `AxChatLogEntry` captures the full prompt sent to the model and its response: + +```typescript +type AxChatLogMessage = + | { role: 'system'; content: string } + | { role: 'user'; content: string } + | { role: 'assistant'; content: string } + | { role: 'tool'; name: string; content: string }; + +type AxChatLogEntry = { + name?: string; // e.g. "distiller", "executor", "responder" + model: string; + messages: AxChatLogMessage[]; + modelUsage?: AxProgramUsage; + stage?: 'ctx' | 'task'; +}; +``` + +### getUsage() + +Returns token usage split by actor/responder. Each sub-array contains one `AxProgramUsage` entry per model/run, merged by `(ai, model)` key. + +```typescript +const usage = myAgent.getUsage(); +// { actor: AxProgramUsage[], responder: AxProgramUsage[] } + +console.log('Actor tokens:', usage.actor[0]?.tokens); +console.log('Responder tokens:', usage.responder[0]?.tokens); +``` + +### getStagedUsage() + +Returns usage split by pipeline stage. The `ctx` stage has the distiller actor only; the `task` stage has the executor actor plus responder. + +```typescript +const staged = myAgent.getStagedUsage(); +console.log(staged.ctx?.actor); +console.log(staged.task.actor); +console.log(staged.task.responder); +``` + +### getTraces() + +Returns Ax program traces for the agent pipeline. Use it when the caller needs trace data rather than chat messages or token summaries. + +```typescript +const traces = myAgent.getTraces(); +``` + +### resetUsage() + +Resets both actor and responder usage at once: + +```typescript +myAgent.resetUsage(); +``` + +Type signatures: + +```typescript +// AxAgent +agent.getChatLog(): readonly AxChatLogEntry[] +agent.getUsage(): { actor: AxProgramUsage[]; responder: AxProgramUsage[] } +agent.getStagedUsage(): { ctx?: AxAgentUsage; task: AxAgentUsage } +agent.getTraces(): AxProgramTrace[] +agent.resetUsage(): void + +// AxGen / AxFlow +gen.getChatLog(): readonly AxChatLogEntry[] +gen.getUsage(): AxProgramUsage[] +``` + +## Do Not Generate + +- Do not add both `debug: true` and `actorTurnCallback` unless the user wants both unstructured prompt/runtime visibility and structured telemetry. +- Do not scrape actor prompts or action logs when a callback exposes the data directly. +- Do not let observability callback failures break the agent run; Ax swallows callback errors for telemetry hooks. +- Do not use DSP-layer `onFunctionCall` when the user wants AxAgent runtime function calls. +- Do not enable `showThoughts` unless the user needs provider thought diagnostics and the provider supports it. diff --git a/.claude/skills/ax-agent-optimize/SKILL.md b/.claude/skills/ax-agent-optimize/SKILL.md new file mode 100644 index 0000000..6709ad4 --- /dev/null +++ b/.claude/skills/ax-agent-optimize/SKILL.md @@ -0,0 +1,360 @@ +--- +name: ax-agent-optimize +description: This skill helps an LLM generate correct AxAgent tuning and evaluation code using @ax-llm/ax. Use when the user asks about agent.optimize(...), judgeOptions, eval datasets, optimization targets, saved optimizedProgram artifacts, or agent optimization guidance. +version: "23.0.0" +--- + +# AxAgent Optimize Codegen Rules (@ax-llm/ax) + +Use this skill for `agent.optimize(...)` workflows. Prefer short, modern, copyable patterns. Do not repeat general agent-authoring guidance unless the user needs it. For generic `ax(...)` or `flow(...)` tuning with top-level `optimize(...)`, use the `ax-gepa` skill instead. + +Your job is to help the model choose a good optimization setup for the user's actual goal: + +- If the user wants better tool use, prefer action-aware tasks and either a deterministic metric or the built-in judge depending on how objective the scoring is. +- If the user wants better wording only, responder optimization may be enough. +- If the user wants reusable improvements, include artifact save/load. +- If the user wants cost, tool-use, or child-agent delegation behavior improved, make the eval tasks expose those tradeoffs explicitly. + +## Use These Defaults + +- Use `agent.optimize(...)` only after the agent is already configured and runnable. +- Prefer the built-in judge path first for normal agent tuning. Most users should start with tasks that include `input` and `criteria`, then let `agent.optimize(...)` use its default actor target and judge-based metric. +- Keep top-level `optimize(program, train, metric, options)` for non-agent generators and flows; do not rewrite normal agent task-record examples to the generic helper. +- Prefer a deterministic custom `metric` only when success is easy to score from the prediction and task record. +- Add `judgeAI` plus `judgeOptions` when the judge should run on a stronger or separate model than the agent runtime model. +- Only reach for a plain typed `AxGen` evaluator when the user needs LLM-as-judge behavior outside the built-in `agent.optimize(...)` flow. +- Default optimize target is the actor path; do not surface `target` unless the user clearly wants responder-only tuning or explicit program IDs. +- Use eval-safe tools or in-memory mocks because optimization replays tasks many times. +- Prefer precise tool return schemas such as `f.object(...)` over vague `f.json(...)` whenever the agent must reason about returned fields. +- Prefer task wording with canonical entity names like "the Atlas project" instead of ambiguous labels like "Atlas" when ambiguity could trigger pointless clarification. +- Save artifacts with `axSerializeOptimizedProgram(result.optimizedProgram!)`, then restore with `axDeserializeOptimizedProgram(saved)` and `agent.applyOptimization(...)`. +- For browser-safe persistence, let the caller store the serialized JSON anywhere they want such as localStorage, IndexedDB, or a backend. +- If `bootstrap` is enabled, bootstrapped demos are persisted inside `result.optimizedProgram.demos`; raw failed traces are not saved in v1. +- Auto-promoted context fields (large undeclared inputs kept runtime-only by `autoUpgrade`) appear in captured traces/demos as their truncated preview string, not the full value — same as declared truncate-style `contextFields`. This is expected; do not treat the shortened value as a bug in the saved demos. +- For first examples, pass a plain task array instead of splitting into `train` and `validation` unless the user already has a holdout set. +- GEPA-backed `agent.optimize(...)` now optimizes generic components exposed by the selected target programs; `target: 'actor'` only tunes actor components, `target: 'responder'` only tunes responder components, and `target: 'all'` broadens the component set. +- `result.optimizedProgram.componentMap` is the canonical saved artifact for agent GEPA runs. It may include actor instructions, descriptions, tool descriptions/names, templates, or runtime primitives depending on what the selected target exposes. +- When child-agent delegation matters, expose the child agents as named functions and tune against realistic call/no-call tasks. + +## Decision Guide + +Pick the optimization shape from the user's need: + +- "Make the agent use tools correctly" -> keep the default actor target and use `expectedActions` and `forbiddenActions`. +- "Make final answers read better" -> consider `target: 'responder'`, but only if the task is not mostly tool-selection or clarification behavior. +- "Make the whole agent better" -> use the default actor target first; only broaden target selection when the user clearly wants that extra scope. +- "Tune child-agent delegation" -> use tasks that exercise when to call the child agent, when to call normal tools, and when to answer directly. +- "Compare before and after" -> include a held-out task plus artifact save/load and replay. + +Choose task design carefully: + +- Prefer a small number of realistic tasks over broad but vague datasets. +- Prefer concrete criteria over generic "be helpful" language. +- Prefer explicit action expectations when correctness depends on tools, recipients, dates, or side effects. +- Prefer eval-safe mocks anytime the task touches email, scheduling, external APIs, or persistence. + +## Make Agents Optimizable + +Optimization works much better when the agent and dataset remove avoidable ambiguity: + +- Prefer typed tool outputs over free-form JSON blobs so the actor can rely on exact field names. +- Tell the actor the exact tool fields it may use when payload shape matters. +- Explicitly ban invented fields if the model has any reason to guess hidden IDs or alternate key names. +- If a child agent needs parent values, declare those fields in the child signature and pass them explicitly at the call site. +- For specialist synthesis, tell the agent what narrowed context should be passed to the child agent. +- Keep `maxSubAgentCalls` small in examples unless the user is explicitly testing broad fan-out behavior. +- Use canonical, unambiguous task wording so the model does not burn turns asking for fake clarification. +- In JS-runtime agents, require raw runnable JavaScript only. Ban `javascript:` prefixes, mixed prose/code, and multi-snippet turns. + +Good pattern: + +- tool schema says exactly what fields exist +- task names the exact entity to look up +- actor prompt says which fields to extract before calling a child agent +- metric or judge penalizes unnecessary child-agent calls and tool misuse + +Bad pattern: + +- tool returns `json` with an underspecified shape +- task uses overloaded names like `Atlas` without clarifying whether that is a project, team, or account +- child agent is expected to infer hidden parent state that was never passed in its call arguments +- code agent is allowed to mix natural language with JavaScript in the same turn + +## Metric vs Judge + +Choose the scoring path based on how objectively the task can be measured: + +- Use a custom `metric` when you can score success directly from `prediction` and `example`. +- Use the built-in agent judge when success depends on a full-run qualitative review across tool choices, clarifications, and final output. +- Use `judgeOptions.description` to tell the built-in judge what to value most. +- Use helper-based judge code only when the user is not inside `agent.optimize(...)` and still wants LLM judging. + +Quick rules: + +- Tool correctness with exact expected calls or forbidden calls: prefer a deterministic metric first. +- Simple extraction or classification with known correct answers: prefer a deterministic metric. +- Open-ended assistant quality, nuanced clarification behavior, or broad synthesis quality: prefer the built-in judge. +- GEPA or optimizer flows outside agents that still need LLM judging: use a plain typed `AxGen` evaluator. + +Important: + +- A custom `metric` overrides the built-in judge path entirely. +- Do not introduce a dedicated judge abstraction in new examples; prefer a plain typed `AxGen`. +- Do not add both a custom `metric` and judge guidance unless the user explicitly wants two separate scoring systems and understands only the custom metric drives optimization. +- If the user builds a plain `AxGen` judge metric, prefer a numeric `score:number` output over a string tier when possible. It is simpler and less fragile in practice. + +## Canonical Pattern + +```typescript +import { + AxAIGoogleGeminiModel, + AxJSRuntime, + axDefaultOptimizerLogger, + agent, + ai, + f, + fn, + axDeserializeOptimizedProgram, + axSerializeOptimizedProgram, +} from '@ax-llm/ax'; + +const tools = [ + fn('sendEmail') + .namespace('email') + .description('Send an email message') + .arg('to', f.string('Recipient email address')) + .arg('body', f.string('Email body text')) + .returns( + f.object({ + sent: f.boolean('Whether the email was sent'), + to: f.string('Recipient email address'), + }) + ) + .handler(async ({ to }) => ({ sent: true, to })) + .build(), +]; + +const studentAI = ai({ + name: 'google-gemini', + apiKey: process.env.GOOGLE_APIKEY!, + config: { model: AxAIGoogleGeminiModel.Gemini31FlashLite, temperature: 0.2 }, +}); + +const judgeAI = ai({ + name: 'google-gemini', + apiKey: process.env.GOOGLE_APIKEY!, + config: { model: AxAIGoogleGeminiModel.Gemini35Flash, temperature: 1.0 }, +}); + +const assistant = agent('query:string -> answer:string', { + ai: studentAI, + judgeAI, + contextFields: [], + runtime: new AxJSRuntime(), + functions: tools, + contextPolicy: { preset: 'checkpointed', budget: 'balanced' }, + judgeOptions: { + description: 'Prefer correct tool use over polished wording.', + model: 'judge-model', + }, +}); + +const tasks = [ + { + input: { query: 'Send an email to Jim saying good morning.' }, + criteria: 'Use the email tool and send the message to Jim.', + expectedActions: ['email.sendEmail'], + }, +]; + +const result = await assistant.optimize(tasks, { + maxMetricCalls: 12, + verbose: true, +}); + +const saved = axSerializeOptimizedProgram(result.optimizedProgram!); +const restored = axDeserializeOptimizedProgram(saved); +assistant.applyOptimization(restored); +``` + +## Minimal Normal-User Pattern + +Start here unless the user clearly needs a hand-built scorer: + +```typescript +const tasks = [ + { + input: { query: 'Send an email to Jim saying good morning.' }, + criteria: 'Use the email tool and send the message to Jim.', + expectedActions: ['email.sendEmail'], + }, +]; + +const result = await assistant.optimize(tasks); +assistant.applyOptimization(result.optimizedProgram!); +``` + +- `target` defaults to actor optimization. +- `metric` defaults to the built-in LLM judge. +- `judgeAI` is optional; if omitted, the agent falls back to its configured judge model or runtime model. +- `bootstrap: true` is a good next step for tool-heavy agents when you want GEPA to start from successful traces from the provided tasks. +- The one thing users still need is realistic task records with clear `criteria`. + +## Deterministic Metric Pattern + +Use this when the task has crisp correctness and cost/behavior tradeoffs: + +```typescript +const result = await assistant.optimize(tasks, { + target: 'actor', + metric: ({ prediction, example }) => { + if (prediction.completionType !== 'final' || !prediction.output) { + return 0; + } + + let score = 0; + + if (prediction.output.answer.includes('Jim')) score += 0.4; + + if ( + prediction.functionCalls.some( + (call) => call.qualifiedName === 'email.sendEmail' + ) + ) { + score += 0.4; + } + + if (prediction.turnCount <= 3) { + score += 0.2; + } + + return score; + }, +}); +``` + +Use this pattern when: + +- the task has a known correct answer or exact action pattern +- tool count, child-agent calls, or turn count must be measured explicitly +- you want repeatable, low-variance optimization runs + +## Built-In Judge Pattern + +Use this when the agent behavior needs holistic review: + +```typescript +const result = await assistant.optimize(tasks, { + judgeAI, + judgeOptions: { + model: AxAIGoogleGeminiModel.Gemini35Flash, + description: + 'Be strict about unnecessary child-agent calls, weak clarifications, and incorrect tool choices.', + }, + maxMetricCalls: 12, +}); +``` + +Use this pattern when: + +- task quality is open-ended or hard to score exactly +- the final answer quality matters together with the action trace +- the user wants a judge to consider clarifications, tool errors, and overall completion quality + +## Plain `AxGen` Judge Pattern + +Use this only when the user needs LLM judging outside the built-in `agent.optimize(...)` path: + +```typescript +import { AxGen, s } from '@ax-llm/ax'; + +const judgeGen = new AxGen( + s(` + taskInput:json "Task input", + candidateOutput:json "Candidate output", + expectedOutput?:json "Optional reference output" + -> + score:number "Normalized score from 0 to 1" + `) +); +judgeGen.setInstruction( + 'Score the candidate output from 0 to 1. Reward correctness and task completion. Return only the score field.' +); + +const metric = async ({ prediction, example }) => { + const result = await judgeGen.forward(judgeAI, { + taskInput: example, + candidateOutput: prediction, + expectedOutput: example.expectedOutput, + }); + + return Math.max(0, Math.min(1, result.score)); +}; + +const result = await optimizer.compile(program, train, metric, { + validationExamples: validation, +}); +``` + +Use this pattern when: + +- the user is optimizing an `AxGen`, flow, or another program directly +- the user wants LLM judging without the higher-level `agent.optimize(...)` wrapper +- the user wants to inspect judge results directly, not just a numeric score + +## Dataset And Judge Rules + +- Pass already-loaded tasks. Do not invent a benchmark loader unless the user asks for one. +- Use `expectedActions` and `forbiddenActions` when tool correctness matters. +- `judgeOptions` mirrors normal forward options and supports extra judge guidance through `description`. +- The built-in judge scores from the full agent run, not just the final reply. It can see completion type, clarification payload, final output, action log, normalized function calls, tool errors, and turn count. +- If the user provides a custom `metric`, that overrides the built-in judge path. +- If the user provides an LLM-based custom metric, keep the output schema as small as possible and prefer a direct numeric score. + +Decision rules: + +- Prefer a custom metric when the user has deterministic business scoring, exact action expectations, or explicit cost tradeoffs. +- Prefer the built-in judge when the user wants practical assistant-quality tuning and does not already have a trusted metric. +- Prefer a plain typed `AxGen` evaluator when the user is not calling `agent.optimize(...)` but still wants LLM judging. +- Prefer `judgeOptions.description` to steer the judge toward the user's real priority, such as tool correctness, brevity, groundedness, or policy compliance. + +## Eval Semantics + +- `agent.optimize(...)` runs each evaluation rollout from a clean continuation state. +- Saved runtime state from `getState()` and `setState(...)` is not used during eval rollouts. +- During optimize/eval, `askClarification(...)` is treated as a scored evaluation outcome instead of going through the responder. +- For clarification outcomes in custom metrics, expect `prediction.completionType === 'askClarification'`, populated `prediction.clarification`, and absent `prediction.output`. +- For final outcomes in custom metrics, expect `prediction.completionType === 'final'` and populated `prediction.output`. +- `target: 'responder'` still works, but clarification-heavy tasks are usually low-signal for responder optimization. + +## Delegation Optimization Notes + +- Prefer explicit child agents in `functions: [...]` for specialist delegation. Their calls appear as normal function-call records. +- When delegation behavior matters, tune against the same child-agent/tool structure you expect in production. +- Tell the actor which fields to pass to the child agent and which tasks should stay local. +- For synthesis-style tasks, specify the desired delegation pattern explicitly, for example "call `team.writer(...)` only after narrowing tool output in JS." +- Penalize unnecessary child-agent calls directly in the metric or judge prompt. +- If one training task keeps collapsing to zero, inspect that task first instead of adding more optimizer rounds. Most failures come from task ambiguity, weak tool schemas, or vague delegation guidance rather than GEPA itself. + +## Artifacts And Replay + +- Save `result.optimizedProgram` if the user wants portable artifacts. +- Restore artifacts with `new AxOptimizedProgramImpl(...)`, then call `agent.applyOptimization(...)`. +- Preserve the full optimized program when saving GEPA artifacts; `componentMap` reapplies the learned strings. +- For demonstrations, use fresh eval-safe tool state for baseline, optimize, and restored replay so side effects do not leak across phases. +- If the user wants to show improvement, run a held-out task before optimization, then replay it on a freshly restored optimized agent. + +## Examples + +- [RLM Agent Optimize](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/rlm-agent-optimize.ts) — Gemini office-assistant tuning with save/load +- [AxAgent GEPA Component Optimization](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/axagent-gepa-optimization.ts) — compact support-agent GEPA run with deterministic metric and artifact replay + +## Do Not Generate + +- Do not optimize against production tools with real side effects unless the user explicitly wants that. +- Do not recommend responder-only optimization by default for clarification-heavy workflows. +- Do not omit artifact save/load steps when the user asks for reusable optimized configurations. +- Do not introduce a dedicated judge class or helper abstraction in new agent-optimize examples; prefer the built-in judge path or a plain typed `AxGen`. +- Do not rely on vague `json` tool returns when the agent must reason about specific fields across tool or child-agent calls. +- Do not leave child-agent inputs implicit. If the child needs a fact, pass it explicitly. +- Do not let code-generation agents mix prose and JavaScript if the user is optimizing runtime behavior. diff --git a/.claude/skills/ax-agent-rlm/SKILL.md b/.claude/skills/ax-agent-rlm/SKILL.md new file mode 100644 index 0000000..3602106 --- /dev/null +++ b/.claude/skills/ax-agent-rlm/SKILL.md @@ -0,0 +1,501 @@ +--- +name: ax-agent-rlm +description: This skill helps an LLM generate correct AxAgent RLM/runtime code using @ax-llm/ax. Use when the user asks about RLM code execution, AxJSRuntime, contextFields, contextPolicy, liveRuntimeState, promptLevel, stage prompt controls, executorModelPolicy, maxRuntimeChars, agent.test(...), llmQuery(...), recursionOptions, or long-running agent runtime behavior. +version: "23.0.0" +--- + +# AxAgent RLM Runtime Rules (@ax-llm/ax) + +Use this skill for code-runtime agents and `llmQuery(...)` semantic-helper behavior. For ordinary agent setup, child agents, tool namespaces, clarification, and `bubbleErrors`, use `ax-agent`. For callbacks and logs, use `ax-agent-observability`. For memories and skill loading, use `ax-agent-memory-skills`. + +## Use These Defaults + +- Use `agent(...)`, not `new AxAgent(...)`. +- In stdout-mode RLM, use one observable `console.log(...)` step per non-final actor turn. +- Rely on `autoUpgrade` (ON by default) for oversized inputs you did not declare in `contextFields`: any input value over ~8k serialized chars is kept runtime-only automatically, with a 1,200-char prompt preview plus a `contextMetadata` line, while the full value stays live in the runtime as `inputs.`. Declare a field in `contextFields` only when you want a specific inline policy (`promptMaxChars` / `keepInPromptChars`) or need a large required non-string field kept out of the prompt (those are left inline by auto-upgrade). +- Default to `contextPolicy: { preset: 'checkpointed', budget: 'balanced' }` for most RLM tasks. +- Prefer `contextPolicy: { preset: 'adaptive', budget: 'balanced' }` when older successful turns should collapse sooner while live runtime state stays visible. +- Use `contextMap` for recurring long-context corpora when the distiller should start future runs with a small persisted orientation cache. +- Prefer `promptLevel: 'default'` for normal use. +- Use `promptLevel: 'detailed'` when you want extra anti-pattern examples and tighter teaching scaffolding in the actor prompt. +- Prefer `executorModelPolicy` when the actor may need to upgrade after repeated error turns or discovery in specific namespaces without also upgrading the responder. +- Use explicit child agents in `functions: [...]` when the task needs specialist agents with their own tools/runtime. +- Use `llmQuery(...)` only for focused semantic questions over narrowed context; it does not spawn a tool-using child AxAgent. +- Prefer `maxSubAgentCalls` only when you need an explicit cap on `llmQuery(...)` sub-query usage. + +## Mental Model + +`AxAgent` is a three-stage pipeline. Each `forward()` call walks the stages in order: + +```text +distiller (RLM actor) -> executor (RLM actor) -> responder (synthesizer) +``` + +- **distiller** always runs first. It sees all original inputs so it can understand and normalize the task; declared `contextFields` stay runtime-only when present. It distils relevant evidence by writing runtime-language code in a multi-turn loop, then calls the runtime-exposed `final(request, evidence)` primitive. The request becomes the executor's `inputs.executorRequest`; it must be self-contained and restate the concrete action, target, and constraints, not vague wording like "do it". The distiller should expand the original user task with facts found in context, including follow-ups like "yes, do it". When no `contextFields` are configured, it still performs request normalization over the original inputs with `contextFields: []`. **The distiller has no tools and is not a capability gate.** +- **executor** runs unless the distiller skipped it (below). It receives non-context inputs plus `inputs.executorRequest`, a compact `distilledContextSummary` prompt field, and the real evidence live as `inputs.distilledContext` from the distiller's `final(request, evidence)` payload. Declared or auto-promoted context fields stay runtime-readable as `inputs.` when `contextMetadata` lists them, but their raw contents are not pasted into the executor prompt. The executor owns tool use, decides whether to call its available functions or finish directly from distilled evidence, and reports actual tool results or failures. +- **responder** always runs last. It synthesizes the user's output signature from whichever upstream actor finished the run and must not contradict tool evidence gathered upstream. + +### Direct respond (executor skip) + +With `directResponse: 'auto'` (the default), the distiller can end the run with the `respond(task, evidence)` primitive when the task needs no user-provided functions — the executor stage is skipped entirely (zero executor model calls) and the responder synthesizes straight from the distiller's evidence. Unlike `final`, whose evidence stays live in the shared session by reference, `respond`'s evidence crosses into the responder prompt (budgeted by `maxEvidenceChars`), and the distiller's runtime variables are exported as the cross-run state exactly as the executor's would have been. + +- **Static agents** (no `functions`, no child agents) run respond-only: `final` is not offered to the distiller and every run is distiller → responder. +- **Agents with functions** get `respond` alongside `final` under a conservative covenant: only for tasks answered purely by reading/synthesizing provided context, never when a listed function/module domain covers the need, never for current/live/fresh-state asks (context may be stale — tools are the source of truth for "now"), never for side effects. Landing-gate eval (both pinned models, 3 repeats): 0 false skips on tool-required tasks including a stale-context trap, 100% skip recall on pure context Q&A. +- `directResponse: 'off'` removes the primitive from the prompt and the runtime, and the pipeline rejects a respond payload outright. + +Treat both actor stages as long-running code runtime sessions that the actor steers over multiple turns, not as fresh script generators on every turn. `AxJSRuntime` is the default; custom runtimes set `language` so the actor code field becomes `Code` such as `pythonCode` while JavaScript keeps the legacy `javascriptCode`. + +- Successful code leaves variables, functions, imports, and computed values available in the runtime session. +- The actor should continue from existing runtime state instead of recreating prior work. +- `actionLog`, `liveRuntimeState`, and checkpoint summaries only control what the actor can see again in the prompt. +- Rebuild state only after an explicit runtime restart notice or when you intentionally need to overwrite a value. + +## RLM Actor Code Rules + +Use these rules when generating actor JavaScript for RLM in `AxJSRuntime` stdout mode. For custom runtimes, follow the runtime's `getUsageInstructions()`, primitive overrides, and callable formatter instead. + +- Treat each actor turn as exactly one observable step. +- Inspect what already exists before recomputing it. If a prior turn successfully created a value, prefer reusing that runtime value. +- If you need to inspect a value, compute it or read it, `console.log(...)` it, and stop immediately after that `console.log(...)`. +- On the next turn, continue from the existing runtime state and use the logged result from `Action Log` only as evidence for what happened. +- If the prompt contains `Live Runtime State`, treat it as the canonical view of current variables. +- Errors from child-agent or tool calls appear in `Action Log`; inspect them and fix the code on the next turn. +- Non-final turns should contain exactly one `console.log(...)`. +- Final turns should call `await final(outputGenerationTask, context)` or `await askClarification(...)` without `console.log(...)`. +- Do not write a complete multi-step program in one actor turn. +- Do not combine `console.log(...)` with `await final(...)` or `await askClarification(...)` in the same actor turn. +- Inside actor-authored JavaScript, `await final(...)` and `await askClarification(...)` end the current turn immediately; code after them is dead code. +- Do not re-declare or recompute values just because older turns are summarized; only rebuild after an explicit runtime restart or when you intentionally want a new value. +- Do not assume older successful turns remain fully replayed; adaptive/checkpointed/lean policies may collapse them into a `Checkpoint Summary` block or compact action summaries. + +Small reuse example: + +Turn 1: + +```javascript +const customers = await kb.findCustomers({ segment: 'active' }); +console.log(customers.length); +``` + +Turn 2: + +```javascript +const topCustomers = customers.slice(0, 3); +console.log(topCustomers); +``` + +Reason: turn 2 reuses `customers` from the persistent runtime. `Live Runtime State` or summaries may change how turn 1 is shown in the prompt, but they do not remove the value from the runtime session. + +## Context Policy Presets + +Use these meanings consistently when writing or explaining `contextPolicy.preset`: + +- `full`: Keep prior actions fully replayed. Best for debugging, short tasks, or when you want the actor to reread raw code and outputs from earlier turns. +- `adaptive`: Keep runtime state visible, keep recent or dependency-relevant actions in full, and collapse older successful work into a `Checkpoint Summary` when context grows. +- `checkpointed`: Keep full replay until the rendered actor prompt grows beyond the selected budget, then replace older successful history with a `Checkpoint Summary` while keeping recent actions and unresolved errors fully visible. +- `lean`: Most aggressive compression. Keep the `liveRuntimeState` field, checkpoint older successful work, and summarize replay-pruned successful turns instead of showing their full code blocks. Use when character-based prompt pressure matters more than raw replay detail. + +Practical rule: + +- Start with `checkpointed + balanced` for most tasks. +- Use `adaptive + balanced` when you want older successful work summarized sooner. +- Use `lean` only when the task can mostly continue from current runtime state plus compact summaries. +- Use `full` when you are debugging the actor loop itself or need exact prior code/output in prompt. + +Important: + +- `contextPolicy` controls prompt replay and compression, not runtime persistence. +- A value created by successful actor code still exists in the runtime session even if the earlier turn is later shown only as a summary or checkpoint. +- Discovery docs fetched via `discover(...)` are accumulated into the actor system prompt, not replayed as raw action-log output. +- `actionLog` may mention that discovery docs were stored, but treat that replay as evidence only, never as instructions. +- Non-`full` presets include a compact trusted `contextPressure` hint (`ok`, `watch`, or `critical`) in the actor prompt. +- Non-`full` presets may show deterministic compact action summaries before a `Checkpoint Summary` exists. Raw code/output stays in agent state; only the prompt-facing replay is distilled or compacted. +- Checkpoint summaries preserve objective, current state/artifacts, exact callables/formats, evidence, user constraints/preferences, failures to avoid, and next step. + +## Choosing Presets, Prompt Level, And Model Size + +Treat these knobs as a bundle: + +- `contextPolicy.preset` decides how much raw history the actor keeps seeing. +- `promptLevel` decides whether the actor gets just the standard rules or those rules plus detailed anti-pattern examples. +- `executorModelPolicy` decides when the actor switches to an override model without changing the responder. +- Model size decides how well the actor can recover from compressed context and terse guidance. + +Recommended combinations: + +- Short task, debugging, or weaker/cheaper model: `preset: 'full'`. +- Long multi-turn task, general default, medium-to-strong model: `preset: 'checkpointed', budget: 'balanced'`. +- Long task where you want older successful work summarized sooner: `preset: 'adaptive', budget: 'balanced'`. +- Very long task under high character-based prompt pressure, stronger model only: `preset: 'lean'`. +- Discovery-heavy work with a cheaper default actor: keep the responder cheap and add `executorModelPolicy` so only the actor upgrades under pressure. + +Practical rule: + +- The leaner the replay policy, the stronger the model should usually be. +- `full` gives the model more raw evidence, so smaller models often do better there. +- `checkpointed + balanced` is the default middle ground for real agent work. +- `adaptive + balanced` is the proactive-summarization variant when you want older successful work compressed sooner. +- `lean` should be reserved for models that can reason well from runtime state plus summaries instead of exact old code/output. +- `executorModelPolicy` is usually better than globally upgrading the whole agent when the bottleneck is actor exploration rather than responder synthesis. + +## Option Layout + +Use these top-level controls consistently: + +- `recursionOptions.ai`: routes `llmQuery(...)` sub-query calls to a different AI service than the parent run. +- `recursionOptions.model`, `modelConfig`, and other forward options: tune the AxGen call used by `llmQuery(...)`. +- `maxSubAgentCalls`: shared `llmQuery(...)` sub-query budget across the whole run. Default is `100`. +- `maxBatchedLlmQueryConcurrency`: caps batched `llmQuery([...])` concurrency. +- `maxRuntimeChars`: runtime/output truncation ceiling for console logs, tool results, and interpreter output replay. The effective limit is computed dynamically each turn based on remaining context budget. +- `summarizerOptions`: default model/options for the internal checkpoint summarizer. +- `contextPolicy`: replay/checkpointing/compression policy. +- `contextMap`: optional persistent orientation cache injected into the distiller and updated once after each successful run. `AxAgentContextMap` evolves indefinitely by default; use `{ infiniteEvolve: false, evolveSteps: N }` on the map object for finite warmup followed by reuse. +- `contextOptions`: distiller-stage forward options. +- `autoUpgrade`: smart defaults, ON by default. Auto-enables `functionDiscovery` for large tool catalogs and keeps oversized undeclared input values runtime-only with a truncated prompt preview. Set `false` to opt out, or tune per side: `{ functionDiscovery?: boolean | { aboveFunctionDocChars }, contextFields?: boolean | { promoteAboveChars, previewChars } }`. Explicit `functionDiscovery` and declared `contextFields` always win. +- `executorOptions`: executor-stage forward options such as `description`, `model`, `modelConfig`, `thinkingTokenBudget`, and `showThoughts`. +- `executorModelPolicy`: executor-only model override rules based on consecutive error turns or discovery fetches from listed namespaces. +- `responderOptions`: responder-stage forward options. +- `judgeOptions`: built-in judge options for `agent.optimize(...)`; for tuning workflows use `ax-agent-optimize`. + +Canonical shape: + +```typescript +const researchAgent = agent('query:string -> answer:string', { + contextFields: ['query'], + runtime, + recursionOptions: { + model: 'gpt-5.4-mini', + }, + maxRuntimeChars: 3000, + summarizerOptions: { + model: 'gpt-5.4-mini', + modelConfig: { temperature: 0.1, maxTokens: 180 }, + }, + contextPolicy: { + preset: 'checkpointed', + budget: 'balanced', + }, + contextOptions: { + model: 'gpt-5.4-mini', + maxTurns: 3, + }, + executorOptions: { + description: 'Use tools first and keep JS steps small.', + model: 'gpt-5.4-mini', + }, + executorModelPolicy: [ + { + model: 'gpt-5.4', + aboveErrorTurns: 2, + namespaces: ['db', 'kb'], + }, + ], + responderOptions: { + model: 'gpt-5.4-mini', + }, +}); +``` + +Semantics: + +- `maxRuntimeChars` sets the truncation ceiling and is separate from `contextPolicy.budget`. +- `summarizerOptions` tunes only the internal checkpoint summarizer. It does not change actor or responder model selection. +- `executorModelPolicy` only switches the actor model. It does not change `responderOptions.model`. +- `llmQuery(...)` uses `recursionOptions.ai` when set, otherwise it falls back to the parent `.forward(ai, ...)` service. +- `recursionOptions` configures the AxGen semantic sub-query used by `llmQuery(...)`; it does not create a child AxAgent and cannot give the sub-query tools. +- `executorModelPolicy` entries are ordered from weaker to stronger. If multiple rules match, the last matching entry wins. +- If one entry defines `namespaces`, any successful `discover(...)` function-definition fetch from one of those namespaces marks the rule as matched starting on the next actor turn. +- Do not add `recursionOptions` unless the user needs different model/options for `llmQuery(...)`. + +## Dynamic Output Truncation + +Runtime output truncation is budget-proportional and type-aware: + +- Early turns with little action-log pressure use the full `maxRuntimeChars` ceiling. +- As the action log fills toward `targetPromptChars`, the limit decays linearly down to 15% of the ceiling, hard-floored at 400 chars. +- Large arrays keep the first 3 and last 2 items, with the middle replaced by `... [N hidden items]`. +- Deep objects replace nested values beyond depth 3 with `[Object]` or `[Array(N)]`. +- Error stack traces keep the first 3 and last 1 stack frames. +- Simple values use standard `JSON.stringify` passthrough. + +Users do not need to configure this behavior. `maxRuntimeChars` sets the upper bound; the dynamic system only reduces it. + +## Stage Prompt Controls + +The pipeline has three peer stage-config bags: `contextOptions` (distiller), `executorOptions` (executor), and `responderOptions` (responder). Each accepts the same shape: `description`, `model`, `modelConfig`, `excludeFields`, plus other forward options. + +Key fields: + +- `contextOptions.description`: append extra distiller-specific instructions. +- `executorOptions.description`: append extra executor-specific instructions; this is the typical place for tool-use guidance. +- `responderOptions.description`: append extra responder-specific instructions. +- `contextOptions.model` / `executorOptions.model` / `responderOptions.model`: split model choice across stages. +- `contextOptions.ai` / `executorOptions.ai` / `responderOptions.ai`: override the AI service for a specific stage. +- `executorModelPolicy`: auto-switch only the executor when the run is on a consecutive error streak or discovery fetches land in specific namespaces. + +Good split-model pattern: + +```typescript +const researchAgent = agent('query:string -> answer:string', { + contextFields: ['query'], + runtime, + contextPolicy: { preset: 'checkpointed', budget: 'balanced' }, + executorOptions: { + model: 'gpt-5.4', + }, + responderOptions: { + model: 'gpt-5.4-mini', + }, +}); +``` + +Model guidance: + +- Put the stronger model on the actor when the task depends on multi-turn exploration, discovery, runtime state reuse, or compressed replay. +- Put the stronger model on the responder only when the hard part is final synthesis/formatting rather than exploration. +- For cost-sensitive setups, a common pattern is stronger actor plus cheaper responder. +- Prefer `executorModelPolicy` over globally upgrading the whole agent when the actor only needs help after context grows or the run starts thrashing. + +Prompt/cache shape: + +- Actor turns are compact observable turns, not replayed chat transcripts. +- Stable system prompt: role/stage rules, primitive descriptions, static module list, always-included callable signatures, output contract, and field definitions. +- Cached working inputs: task inputs, inline context, `contextMetadata`, `contextMap`, `memories`, `executorRequest`, `distilledContextSummary`, `discoveredToolDocs`, `loadedSkills`, and `summarizedActorLog`. +- Dynamic turn tail: `guidanceLog`, `actionLog`, `liveRuntimeState`, and `contextPressure`. +- Prefer one compact inspection per non-final turn. Never combine inspection output with `final(...)` or `askClarification(...)`. + +Invalid actor turn: + +```javascript +await discover(['kb.findSnippets']); +const snippets = await kb.findSnippets({ topic: 'severity' }); +await final("Summarize severity findings", { snippets }); +``` + +Reason: this mixes observation and follow-up work in one turn. `discover(...)` returns `void`; read the next prompt's "Discovered Tool Docs" section before calling the function. + +## AxJSRuntime Security + +Default `new AxJSRuntime()` is hardened: no network, no filesystem, no child process, dynamic `import()` blocked, intrinsics frozen, `ShadowRealm` locked to `undefined`, worker IPC locked in browser/Deno/Bun, Bun workers use `smol: true`, and on Node 20+ the OS Permission Model auto-engages where available. + +Threat model: this is defense-in-depth for LLM-authored code, not a container or VM boundary. Host callbacks and granted runtime permissions remain the authority boundary; keep durable secrets and privileged effects in host-side functions. + +Permission enum (`AxJSRuntimePermission`): +`NETWORK`, `STORAGE`, `CODE_LOADING`, `COMMUNICATION`, `TIMING`, `WORKERS`, `FILESYSTEM`, `CHILD_PROCESS`. + +Options quick reference: + +- `permissions?: readonly AxJSRuntimePermission[]`: default `[]`; opt in capabilities. +- `blockDynamicImport?: boolean`: default `true`. +- `allowedModules?: readonly string[]`: default `[]`; narrow dynamic-import allowlist gate. Allowlisted specifiers are attempted, but full Node module namespace passthrough depends on Node vm semantics. +- `freezeIntrinsics?: boolean`: default `true`. +- `blockShadowRealm?: boolean`: default `true`. +- `lockWorkerIPC?: boolean`: default `true`. +- `preventGlobalThisExtensions?: boolean`: default `false`; opt-in and breaks top-level persistence. +- `useNodePermissionModel?: boolean | 'auto'`: default `'auto'`. +- `nodePermissionAllowlist?: { fsRead?; fsWrite?; childProcess?; addons?; wasi? }`. +- `resourceLimits?: { maxOldGenerationSizeMb?; maxYoungGenerationSizeMb?; codeRangeSizeMb?; stackSizeMb? }`. +- `allowDenoRemoteImport?: boolean`: default `false`. +- `allowUnsafeNodeHostAccess?: boolean`: default `false`. + +Recipes: + +```typescript +new AxJSRuntime(); + +new AxJSRuntime({ permissions: [AxJSRuntimePermission.NETWORK] }); + +new AxJSRuntime({ + permissions: [AxJSRuntimePermission.FILESYSTEM], + allowedModules: ['node:fs', 'node:fs/promises', 'node:path'], + useNodePermissionModel: 'auto', + nodePermissionAllowlist: { + fsRead: ['/app/data'], + fsWrite: ['/app/data'], + }, +}); +``` + +Rules for the LLM author: + +- Default to `new AxJSRuntime()` with no options unless the user asked for a specific capability. +- When the user asks for `fetch`, add `permissions: [AxJSRuntimePermission.NETWORK]`. +- When the user asks for filesystem access, prefer host-side tool functions. If direct runtime filesystem access is required, add `permissions: [AxJSRuntimePermission.FILESYSTEM]`, scope with `nodePermissionAllowlist` when the user names a directory, and treat `allowedModules` as an import allowlist gate rather than a portability guarantee. +- Do not disable `freezeIntrinsics`, `blockShadowRealm`, or `lockWorkerIPC` unless the user explicitly asks. +- Treat `allowUnsafeNodeHostAccess: true` as a red flag; only use it when the user is authoring trusted code in their own process. +- `preventGlobalThisExtensions: true` breaks top-level `var`/`let`/`const` persistence across turns; never set it for stdout-mode RLM where persistence is load-bearing. +- On Deno, `blockDynamicImport` is a no-op; the defense is the worker permission sandbox. Pass `allowDenoRemoteImport: true` only if remote module loading is genuinely required. + +## Custom Code Runtimes + +Implement `AxCodeRuntime` when the actor should write a language other than JavaScript. + +- Set `language` to the model-facing language name. JavaScript aliases (`JavaScript`, `js`, `ecmascript`) keep `javascriptCode`; other values derive lower-camel code fields such as `pythonCode` or `cSharpCode`. +- Keep execution inside `createSession(globals, options)`. AxAgent passes `inputs`, `llmQuery`, `final`, `askClarification`, progress callbacks, memory/discovery primitives, and namespaced tools as host globals; the runtime decides how those globals appear in the target language. +- Put language syntax, output behavior, persistence semantics, and completion-call examples in `getUsageInstructions()`. +- Use `getPrimitiveOverrides()` to describe language-native calls for built-in primitives, and `formatCallable()` to describe language-native calls for tools and child agents. +- Implement `inspectGlobals()` on sessions when `contextPolicy` should show live runtime state for non-JavaScript runtimes; otherwise AxAgent will not run JavaScript fallback inspection snippets. + +## RLM Test Harness + +Use `agent.test(code, contextFieldValues?, options?)` when the user wants to validate runtime snippets against the actual AxAgent runtime environment without running the full actor/responder loop. With `AxJSRuntime`, those snippets are JavaScript. + +```typescript +import { AxJSRuntime, agent, f, fn } from '@ax-llm/ax'; + +const runtime = new AxJSRuntime(); + +const tools = [ + fn('sum') + .description('Return the sum of the provided numeric values') + .namespace('math') + .arg('values', f.number('Value to add').array()) + .returns(f.number('Sum of all values')) + .handler(async ({ values }) => + values.reduce((total, value) => total + value, 0) + ) + .build(), +]; + +const toolHarness = agent('query:string -> answer:string', { + contextFields: [], + runtime, + functions: tools, + contextPolicy: { preset: 'checkpointed', budget: 'balanced' }, +}); + +const toolOutput = await toolHarness.test( + 'console.log(await math.sum({ values: [3, 5, 8] }))' +); + +console.log(toolOutput); +``` + +Rules: + +- `test(...)` creates a fresh runtime session per call. +- Context-field snippets run in the context/distiller runtime and expose `inputs` plus non-colliding top-level aliases for configured `contextFields`. +- Tool snippets should use an agent with no `contextFields`, or test the executor stage directly, so namespaced functions, child agents, and `llmQuery(...)` are in scope. +- In `AxJSRuntime`, do not rely on calling `inspectRuntime()` from inside `test(...)` snippets yet; prefer checking runtime globals directly inside the snippet. +- It returns the formatted runtime output string. +- It throws on runtime failures instead of returning LLM-style error strings. +- Do not call `final(...)` or `askClarification(...)` inside `test(...)` snippets. +- Pass only `contextFields` values to `test(...)`; it is not a general way to inject arbitrary non-context inputs. +- If the snippet uses `llmQuery(...)`, provide an AI service through the agent config or `options.ai`. + +## `llmQuery(...)` Rules + +Available forms: + +- `await llmQuery(query, context?)` +- `await llmQuery({ query, context? })` +- `await llmQuery([{ query, context }, ...])` + +Rules: + +- `llmQuery(...)` forwards only the explicit `context` argument. +- Parent inputs, runtime variables, tool results, and discovered docs are not automatically available to `llmQuery(...)`; include any needed facts in `context`. +- `llmQuery(...)` is a direct semantic helper backed by an AxGen sub-query. It does not create a child AxAgent, does not run an actor runtime session, and does not have access to tools or discovery. +- Use batched `llmQuery([...])` only for independent semantic questions. Use serial calls when later work depends on earlier results. +- Pass compact named object context instead of huge raw parent payloads. +- Do not assume anything other than the returned string comes back from `llmQuery(...)`. +- `maxSubAgentCalls` is a shared budget for `llmQuery(...)` sub-queries across the top-level run. +- Single-call `llmQuery(...)` may return `[ERROR] ...` on non-abort failures. +- Batched `llmQuery([...])` returns per-item `[ERROR] ...`. +- If a result starts with `[ERROR]`, inspect or branch on it instead of assuming success. + +Minimal example: + +```javascript +const summary = await llmQuery('Summarize this incident', inputs.context); +if (summary.startsWith('[ERROR]')) { + console.log(summary); +} else { + console.log(summary); +} +``` + +Parallel semantic review example: + +```javascript +const narrowedIncidents = incidents.map((incident) => ({ + id: incident.id, + timeline: incident.timeline, + notes: incident.notes.slice(0, 1200), +})); + +const [severityReview, followupReview] = await llmQuery([ + { + query: + 'Use discovery and available tools to review severity policy alignment. Return compact findings.', + context: { + incidents: narrowedIncidents, + rubric: 'severity-policy', + }, + }, + { + query: + 'Use discovery and available tools to review postmortem and follow-up obligations. Return compact findings.', + context: { + incidents: narrowedIncidents, + rubric: 'postmortem-followup', + }, + }, +]); + +const merged = await llmQuery( + 'Merge these delegated reviews into one manager-ready summary with next steps.', + { + severityReview, + followupReview, + audience: inputs.audience, + } +); +``` + +Delegation decision guide: + +- **JS-only**: deterministic logic such as filter, sort, count, regex, or date math -> do it inline. +- **Single-shot semantic**: needs LLM reasoning but no tools or multi-step exploration -> single `llmQuery(...)` with narrow context. +- **Specialist/tool delegation**: needs its own tools, discovery, runtime, or reusable role -> create a child `agent(...)` and pass it in `functions: [...]`. +- **Parallel semantic fan-out**: two or more independent semantic-only subtasks -> batched `llmQuery([...])`. + +Context handling: + +- Always narrow with JS before delegating. Never pass raw `inputs.*`. +- Name context keys semantically, e.g. `{ emails: filtered, rubric: 'classify-urgency' }`. +- Estimate total sub-query calls before fanning out. `maxSubAgentCalls` is shared across the run. + +Patterns: + +- Fan-Out / Fan-In: JS narrows into categories -> `llmQuery([...])` fans out per category -> JS or one more `llmQuery(...)` merges semantic results. +- Pipeline: serial `llmQuery(...)` calls where each depends on the prior result. +- Specialist tool use: call child agents or tools via their namespaced function globals, e.g. `await team.writer({ draft })`. + +## Examples + +Fetch these for full working code: + +- [RLM](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/rlm.ts) - RLM basic +- [RLM Long Task](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/rlm-long-task.ts) - RLM context policy +- [RLM Discovery](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/rlm-discovery.ts) - discovery mode, grouped tools, child agents as functions, and semantic `llmQuery(...)` +- [RLM Adaptive Replay](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/rlm-adaptive-replay.ts) - adaptive replay + +Flagship real-world long-agents (also ported to Python, Go, Rust, Java, and C++ under `src/examples//long-agents/`; run with `npm run example -- `): + +- [Incident Log Forensics](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/typescript/long-agents/incident-log-forensics.ts) - large-context log forensics over `contextFields` (Gemini) +- [Codebase Peek Map](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/typescript/long-agents/codebase-peek-map.ts) - Peek-paper context-map orientation over a large repo snapshot +- [Data Analyst with Tools](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/typescript/long-agents/data-analyst-with-tools.ts) - large data dictionary in `contextFields` + typed warehouse tools the model queries instead of inlining +- [Smart Defaults Agent](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/typescript/long-agents/smart-defaults-agent.ts) - oversized undeclared context auto-promoted runtime-only, with relevance hints and runtime tools +- [Self-Improving Lab](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/typescript/long-agents/self-improving-lab.ts) - many-tool agent that runs experiments, grades them with an independent verifier, and distills verified rules into memory + +## Do Not Generate + +- Do not write a full multi-step RLM actor program in one turn. +- Do not combine `console.log(...)` with `final(...)`. +- Do not assume old successful turns stay fully replayed under adaptive/checkpointed/lean policies. +- Do not rebuild runtime state just because a prior turn was summarized. +- Do not describe `llmQuery(...)` as spawning a tool-using child AxAgent. +- Do not assume parent inputs are available to `llmQuery(...)` unless passed in `context`. +- Do not ignore `[ERROR] ...` results from `llmQuery(...)`. +- Do not grant `AxJSRuntime` permissions unless the user asked for the capability. diff --git a/.claude/skills/ax-agent/SKILL.md b/.claude/skills/ax-agent/SKILL.md new file mode 100644 index 0000000..abd023c --- /dev/null +++ b/.claude/skills/ax-agent/SKILL.md @@ -0,0 +1,631 @@ +--- +name: ax-agent +description: This skill helps an LLM generate correct core AxAgent code using @ax-llm/ax. Use when the user asks about agent(), child agents, namespaced functions, discovery mode, clarification, bubbleErrors, host-side final/clarification protocol, or ordinary agent runtime behavior. For RLM/code-runtime work use ax-agent-rlm; for callbacks and telemetry use ax-agent-observability; for recall/memory/skill loading use ax-agent-memory-skills; for agent.optimize(...) use ax-agent-optimize. +version: "23.0.0" +--- + +# AxAgent Codegen Rules (@ax-llm/ax) + +Use this skill to generate small, correct `AxAgent` code. Prefer modern factory-style APIs and copyable patterns. Do not write tutorial prose unless the user explicitly asks for explanation. + +Your job is to choose the smallest correct `AxAgent` shape for the user's needs: + +- If the user wants a normal tool-using assistant, keep the config minimal. +- If the user wants long-running code execution, use the `ax-agent-rlm` skill. +- If the user wants callbacks, logs, tracing, or usage data, use the `ax-agent-observability` skill. +- If the user wants dynamic memory retrieval or skill-guide loading, use the `ax-agent-memory-skills` skill. +- If the user wants tuning or eval with `agent.optimize(...)`, use the `ax-agent-optimize` skill. + +## Use These Defaults + +- Use `agent(...)`, not `new AxAgent(...)`. +- Prefer string signatures or `f()` signatures over hand-written signature objects. +- Put `ai`, `judgeAI`, and `agentIdentity` on the `agent(...)` config when you want instance defaults or child-agent metadata. +- Prefer `fn(...)` for host-side function definitions instead of hand-writing JSON Schema objects. +- Prefer namespaced functions such as `utils.search(...)` or `kb.find(...)`. +- Pass child agents directly in `functions: [...]`. They land under their `agentIdentity.namespace` (or `utils` if unset), exactly like a `fn()` tool. +- If discovery is enabled, call `discover(...)` before using callables whose docs are not already in the prompt. +- Use explicit child agents in `functions: [...]` for specialist delegation; do not model that as recursive `llmQuery(...)`. +- Add `bubbleErrors` only for fatal infrastructure errors that should abort `.forward()`. + +## Decision Guide + +Map user intent to agent shape before writing code: + +- "Use tools and answer" -> plain `agent(...)` with local functions, no extra observability. +- "Need child agents with distinct responsibilities" -> add child agents to the parent's `functions: [...]` list and set each child's `agentIdentity.namespace` when you want a specific runtime call site such as `team.writer(...)`. +- "Need tool discovery because names/schemas are not stable" -> enable discovery and generate discovery-first actor code. +- "Need certain errors to escape the agent loop" -> add `bubbleErrors` with error classes; those errors propagate through function handlers, actor code, and `llmQuery(...)` sub-queries to `.forward()`. +- "Inspect large context with code", "RLM", or "`llmQuery(...)`" -> use `ax-agent-rlm`. +- "Need debugging, traces, progress updates, tool-call logs, chat logs, or usage" -> use `ax-agent-observability`. +- "Need memories, recall, dynamic skill guides, `discover({ skills })`, or loaded/used tracking" -> use `ax-agent-memory-skills`. + +## Critical Rules + +- Use `agent(...)` factory syntax for new code. +- Add child agents to the parent's `functions: [...]` list. Each child's `agentIdentity.namespace` (or `utils`, the default) determines the runtime call site, e.g. `await team.writer({...})`. +- If discovery is enabled, call `discover(...)` before using callables whose docs are not already in the prompt. +- `autoUpgrade` is ON by default: large tool catalogs auto-enable discovery, and oversized undeclared input values are auto-kept runtime-only with a truncated prompt preview. Explicit `functionDiscovery` and declared `contextFields` always win; set `autoUpgrade: false` to opt out. +- `directResponse` is ON by default (`'auto'`): when a task needs no user-provided functions, the distiller ends the run with `respond(task, evidence)` and the executor stage is skipped (zero executor model calls). Function-less agents run respond-only every time; agents with functions offer `respond` under a conservative covenant (no live/fresh-state asks, no side effects, nothing a listed function/module domain covers). Set `directResponse: 'off'` to always run the executor. +- If a host-side `AxAgentFunction` needs to end the current actor turn, use `extra.protocol.final(...)` or `extra.protocol.askClarification(...)`. +- In public `forward()` and `streamingForward()` flows, `askClarification(...)` throws `AxAgentClarificationError`; it does not go through the responder. +- When resuming after clarification, prefer `error.getState()` from the thrown `AxAgentClarificationError`, then call `agent.setState(savedState)` before the next `forward(...)`. +- Errors listed in `bubbleErrors` bypass actor-loop catch blocks and propagate directly to the caller of `.forward()`. +- Child agents receive only the arguments the actor passes. Pass parent fields explicitly via `inputs.` or use `inputUpdateCallback` when many calls need the same value. +- Audio input fields are transcribed before agent planner/executor/responder stages by default; internal agent stages receive text transcripts, not base64 audio. + +## Canonical Pattern + +```typescript +import { agent, ai, f } from '@ax-llm/ax'; + +const llm = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, +}); + +const assistant = agent( + f() + .input('query', f.string()) + .output('answer', f.string()) + .build(), + { + agentIdentity: { + name: 'Assistant', + description: 'Answers user questions', + }, + contextFields: [], + } +); + +const result = await assistant.forward(llm, { query: 'What is TypeScript?' }); +console.log(result.answer); +``` + +## Audio Inputs And Speech Outputs + +Agents can accept audio inputs and return scripted speech artifacts. The runtime transcribes audio input fields before internal stages run, then synthesizes `:audio` outputs after the final structured response is selected. + +```typescript +const voiceAgent = agent( + 'recording:audio, question:string -> speech:audio, summary:string', + { + agentIdentity: { + name: 'Voice Assistant', + description: 'Answers spoken requests', + }, + contextFields: [], + } +); + +const result = await voiceAgent.forward( + llm, + { + recording: { data: base64Wav, format: 'wav' }, + question: 'What should I do next?', + }, + { + speech: { + transcribe: { model: 'gpt-4o-mini-transcribe' }, + speak: { voice: 'alloy', format: 'mp3' }, + }, + } +); + +console.log(result.summary); +console.log(result.speech.data); +``` + +Use direct `ax(...)` or `.chat()` if the model should receive native audio instead of a transcript-first agent pipeline. + +## Child Agents As Tools + +Child agents are passed in the parent's `functions` list. There is no separate `agents` option for new code. Each child agent's `agentIdentity.namespace` (or `utils`, the default) determines where it lands in the actor runtime. With `AxJSRuntime`, that produces JavaScript call sites such as `team.writer(...)`: + +```typescript +const writer = agent('draft:string -> revision:string', { + agentIdentity: { + name: 'Writer', + description: 'Polishes drafts', + namespace: 'team', + }, + contextFields: [], +}); + +const coordinator = agent('query:string -> answer:string', { + functions: [writer], + contextFields: [], +}); +``` + +Generated runtime call: + +```javascript +const result = await team.writer({ draft: '...' }); +``` + +Without `agentIdentity.namespace`, the child lands under `utils.` like any other tool: + +```javascript +const result = await utils.writer({ draft: '...' }); +``` + +Rules: + +- Add child agents to `functions: [...]`, the same array as `fn(...)` tools. +- Set `agentIdentity.namespace` on the child to control its runtime call site. +- `onFunctionCall` observers receive `kind: 'internal'` for agent-derived calls and `kind: 'external'` for user-registered tools. + +### Reserved namespace names + +The agent runtime injects a fixed set of globals into the runtime session. These names cannot be used as `agentIdentity.namespace` values or as agent-function namespaces. + +```text +inputs +llmQuery +final +askClarification +reportSuccess +reportFailure +inspectRuntime +discover +recall +``` + +Pick any other lowercase identifier such as `utils`, `kb`, `tools`, `team`, or `db`. + +## Tool Functions And Namespaces + +```typescript +import { agent, f, fn } from '@ax-llm/ax'; + +const findSnippets = fn('findSnippets') + .description('Find handbook snippets by topic') + .namespace('kb') + .arg('topic', f.string('Topic keyword')) + .returns(f.string('Matching snippet').array()) + .example({ + title: 'Find severity guidance', + code: 'await kb.findSnippets({ topic: "severity" });', + }) + .handler(async ({ topic }) => []) + .build(); + +const analyst = agent('query:string -> answer:string', { + functions: [findSnippets], + contextFields: [], +}); +``` + +Generated runtime call: + +```javascript +const snippets = await kb.findSnippets({ topic: 'severity' }); +``` + +Rules: + +- Prefer namespaced functions. +- Default function namespace is `utils` when no namespace is set. +- With `AxJSRuntime`, use the runtime call shape `await .({...})`. Custom runtimes should expose equivalent namespaced calls through their own `formatCallable()` guidance. +- `.arg()` and `.returns()` can use Ax field helpers or any Standard Schema v1 validator directly. + +## Grouped Function Modules + +For discovery mode, group functions into modules using the `AxAgentFunctionGroup` shape when you want a clean namespace tree such as `kb.find(...)` or `metrics.score(...)` without setting `namespace` on every individual `fn(...)`: + +```typescript +const parent = agent('query:string -> answer:string', { + functions: [ + { + namespace: 'kb', + title: 'Knowledge Base', + selectionCriteria: 'Use for handbook and documentation lookups.', + description: 'Knowledge base lookups', + functions: [findSnippetsFn, searchPagesFn], + }, + { + namespace: 'workflow', + title: 'Workflow Controls', + description: 'Small control functions the actor should always see', + alwaysInclude: true, + functions: [completeFn], + }, + ], + functionDiscovery: true, + contextFields: [], +}); +``` + +MCP clients and other `toFunction()` providers can be placed directly inside a group after initialization: + +```typescript +await mcpClient.init(); + +const parent = agent('query:string -> answer:string', { + functions: [ + { + namespace: 'memory', + title: 'Memory MCP', + description: 'Memory server tools', + selectionCriteria: 'Use for persistent memory lookup and updates.', + functions: [mcpClient], + }, + ], + functionDiscovery: true, + contextFields: [], +}); +``` + +Rules: + +- A group is `{ namespace, title, description, functions: [...] }`. +- `selectionCriteria` is optional but useful in discovery mode; it tells the actor when to choose that module. +- The group's `namespace`, `title`, `selectionCriteria`, and `description` show up in `discover(...)` module docs. +- `relevanceRanking` (default ON — set `false` to opt out): a deterministic local ranker that injects an advisory `### Likely Relevant` shortlist into the executor turn (dynamic, non-cached field — the cached prompt stays byte-stable). Enabled by default after its A/B gate passed on both small and frontier models and implemented in the generated language ports through AxIR Core. Details in `ax-agent-memory-skills`; outcomes observable via the `relevance_ranking` context event (`ax-agent-observability`). +- Add `alwaysInclude: true` to a group when discovery mode is on but the actor should always see that group's full callable definitions inline in the prompt. +- Keep `functions: [...]` either flat or grouped. Runtime validation rejects mixed plain function entries and group objects. +- In flat mode, pass `fn(...)` tools, child agents, and `toFunction()` providers directly. +- In grouped mode, put callable entries and `toFunction()` providers inside groups. To expose a child agent inside a group, use `childAgent.getFunction()`. + +## Host-Side Completion From Functions + +Use this pattern when the actor should call a namespaced function, but the host-side function implementation should decide to end the turn: + +```typescript +import { f, fn } from '@ax-llm/ax'; + +const finishReply = fn('finishReply') + .description('Complete the actor turn with the final reply text') + .namespace('workflow') + .arg('reply', f.string('Final reply text')) + .returns(f.string('Final reply text')) + .handler(async ({ reply }, extra) => { + extra?.protocol?.final(reply); + return reply; + }) + .build(); + +const askForOrderId = fn('askForOrderId') + .description('Complete the actor turn by requesting clarification') + .namespace('workflow') + .arg('question', f.string('Clarification question')) + .returns(f.string('Clarification question')) + .handler(async ({ question }, extra) => { + extra?.protocol?.askClarification(question); + return question; + }) + .build(); +``` + +Rules: + +- `extra.protocol` is only available when the function call comes from an active AxAgent actor runtime session. +- Use `extra.protocol.final(...)`, `extra.protocol.askClarification(...)`, or `extra.protocol.guideAgent(...)` only inside host-side function handlers. +- Inside actor-authored runtime code, use the runtime globals `final(...)` and `askClarification(...)` with the syntax documented by the active runtime. +- `extra.protocol.guideAgent(...)` is handler-only internal control flow. It stops the current actor turn and appends trusted guidance to `guidanceLog` for the next iteration. +- `askClarification(...)` accepts either a simple string or a structured object with `question` plus optional UI hints such as `type: 'date' | 'number' | 'single_choice' | 'multiple_choice'` and `choices`. + +## Clarification And Resume State + +Use this pattern when the actor should pause for user input and continue later from the same runtime state. + +```typescript +import { + AxAgentClarificationError, + AxJSRuntime, + agent, + ai, +} from '@ax-llm/ax'; + +const llm = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, +}); + +const tripAgent = agent('request:string, answer?:string -> reply:string', { + contextFields: [], + runtime: new AxJSRuntime(), +}); + +let savedState = tripAgent.getState(); + +try { + await tripAgent.forward(llm, { + request: 'Plan a Lisbon trip', + }); +} catch (error) { + if (error instanceof AxAgentClarificationError) { + console.log(error.question); + savedState = error.getState(); + } else { + throw error; + } +} + +if (savedState) { + tripAgent.setState(savedState); + const resumed = await tripAgent.forward(llm, { + request: 'Plan a Lisbon trip', + answer: 'June 1-5', + }); + console.log(resumed.reply); +} +``` + +Public flow rules: + +- `forward()` and `streamingForward()` throw `AxAgentClarificationError` when the actor calls `askClarification(...)`. +- Successful `final(...)` completions always continue through the responder in public flows. +- `AxAgentClarificationError.question` is the user-facing question text. +- `AxAgentClarificationError.clarification` is the normalized structured payload. +- `AxAgentClarificationError.getState()` returns the saved continuation state captured at throw time. +- `agent.getState()` and `agent.setState(...)` export or restore continuation state on the agent instance. +- `test(...)` is different: it returns structured completion payloads for harness/debug use instead of throwing clarification exceptions. + +Structured clarification payloads: + +- String shorthand is allowed: `askClarification("What dates should I use?")`. +- Structured form is preferred for richer chat UIs: + +```javascript +askClarification({ + question: 'Which route should I use?', + type: 'single_choice', + choices: ['Fastest', 'Scenic'], +}); +``` + +- Supported `type` values are `text`, `number`, `date`, `single_choice`, and `multiple_choice`. +- `single_choice` payloads with missing, empty, or malformed `choices` are downgraded to a plain clarification question instead of failing the turn. +- `multiple_choice` payloads must include at least two valid choices; otherwise the actor turn fails with a corrective runtime error. +- Choice entries may be strings or `{ label, value? }` objects. +- Invalid clarification payloads such as a missing `question` are actor-turn runtime errors, not successful clarification completions. + +State notes: + +- `runtimeBindings` restores execution state; `runtimeEntries`, `actionLogEntries`, and `checkpointState` restore prompt context. +- Resume does not create a fake rehydration action-log turn; provenance still points to the original actor code that set the value. +- Only serializable/structured-clone-friendly values are guaranteed to round-trip through `getState()` / `setState(...)`. +- Reserved runtime globals such as `inputs`, tools, and protocol helpers are rebuilt fresh and are not part of saved state. +- Treat one agent instance as conversation-scoped when using `setState(...)`; do not share one mutable resumed instance across unrelated concurrent conversations. + +## Bubble Errors + +Use `bubbleErrors` when certain exceptions thrown inside function handlers or `llmQuery(...)` sub-query calls should propagate all the way out to `.forward()` instead of being caught by the actor loop and returned as `[ERROR]` strings. + +```typescript +import { agent, f, fn } from '@ax-llm/ax'; + +class DatabaseError extends Error { + constructor(message: string) { + super(message); + this.name = 'DatabaseError'; + } +} + +const dbTool = fn('queryUsers') + .description('Query the user database') + .namespace('db') + .arg('filter', f.string('Filter expression')) + .returns(f.string('JSON result')) + .handler(async ({ filter }) => { + if (!isConnected()) throw new DatabaseError('DB connection refused'); + return JSON.stringify(await db.query(filter)); + }) + .build(); + +const myAgent = agent('query:string -> answer:string', { + contextFields: [], + functions: [dbTool], + bubbleErrors: [DatabaseError], +}); +``` + +Rules: + +- `bubbleErrors` takes an array of Error constructor classes, checked via `instanceof`. +- A matching error thrown inside a function handler, during actor code execution, or inside an `llmQuery(...)` sub-query propagates immediately to `.forward()`. +- Use `bubbleErrors` for fatal infrastructure errors such as DB down, auth failure, or quota exceeded. +- Do not use `bubbleErrors` for expected recoverable errors; let those return as `[ERROR] ...` strings so the actor can handle them. +- `AxAgentClarificationError` and `AxAIServiceAbortedError` always bubble up unconditionally. + +## Unified Final Signal + +There are two ways to end a successful run through the responder: + +1. In actor JS code, call `final(message)` when no extra context object is needed, or `final(task, context)` when you gathered evidence. +2. In function handlers, use `extra.protocol.final(...)` with the same one-arg or two-arg forms. + +Rules: + +- Use `final(message)` when the actor already knows the answer and no extra context object is needed. +- Use `final(task, context)` when context was gathered and needs synthesis into output fields. +- In function handlers, use `extra.protocol.final(...)` instead of a separate respond API. +- The responder still runs for both successful `final(...)` forms. +- Use `askClarification(...)` when the user must provide more information to continue. + +## Discovery Mode + +Enable discovery mode when you want the actor to discover modules and fetch callable definitions on demand: + +```typescript +const analyst = agent('context:string, query:string -> answer:string', { + agentIdentity: { + name: 'Analyst', + description: 'Analyzes long context', + namespace: 'team', + }, + contextFields: ['context'], + functions: [writer, ...tools], + functionDiscovery: true, +}); +``` + +Discovery API: + +- `await discover(item: string): void` +- `await discover(items: string[]): void` +- `await discover({ tools?: string | string[], skills?: string | string[] }): void` when `onSkillsSearch` is configured + +Discovery returns `void`; fetched docs render in the next executor prompt. + +Rules: + +- `discover('kb')` loads a module callable list when `kb` is a discoverable module. +- `discover('kb.findSnippets')` loads a full callable definition. +- `discover('lookup')` resolves as `utils.lookup`. +- `discover({ tools: ['kb'], skills: ['release-checklist'] })` loads tool docs and skill bodies in one turn. +- Call one batched `discover(...)` with every module, callable, and skill you need. +- Do not split discovery into separate calls or wrap discovery in `Promise.all(...)`. +- Read the next prompt's "Discovered Tool Docs" and "Loaded Skills" sections. +- If a guessed call fails, stop guessing nearby names. Run `discover(...)` for that module or function and call only the exact discovered qualified name. + +## Threading Parent Fields Into Child Agents + +If a child agent requires a parent field such as `audience`, declare it on the child's signature and pass it explicitly when calling the child from the actor: + +```typescript +const writingCoach = agent('draft:string, audience:string -> revision:string', { + agentIdentity: { + name: 'Writing Coach', + description: 'Polishes summaries for a target audience', + namespace: 'team', + }, + contextFields: [], +}); + +const analyst = agent('context:string, audience:string, query:string -> answer:string', { + functions: [writingCoach], + contextFields: ['context'], +}); +``` + +Generated runtime call: + +```javascript +const polished = await team.writingCoach({ + draft: summary, + audience: inputs.audience, +}); +``` + +Rules: + +- Pass parent fields explicitly via the call site. +- If many children need the same field on every call, use `inputUpdateCallback` to inject the value before each executor turn. +- Do not assume auto-propagation; child agents receive only the args the actor passes. + +## Core API Reference + +Factory shape: + +```typescript +agent(signature, { + ai, + judgeAI, + agentIdentity, + contextFields, + functions, + functionDiscovery, + autoUpgrade, + ...agentOptions, +}); +``` + +- `ai` is an optional default service for the agent instance; `.forward(ai, ...)` can still pass the runtime service. +- `judgeAI` is the optional default judge/teacher service used by optimize flows. +- `agentIdentity` controls the user-facing agent identity and child-agent function metadata. + +```typescript +agentIdentity?: { + name: string; + description: string; + namespace?: string; +} +``` + +- `name` is normalized to camelCase for child-agent function names. +- `name` and `description` are included in the actor and responder prompts as the user-facing agent identity. +- `namespace` changes the child-agent module from default `utils` to a custom module such as `team`. + +Each `contextFields` entry is either a plain field name string or an object controlling how much of the value is inlined into the distiller prompt: + +- `{ field, promptMaxChars: N }`: inline only when the serialized value is at most `N` chars; otherwise omit it from the prompt and keep it runtime-only. +- `{ field, keepInPromptChars: N, reverseTruncate?: boolean }`: always inline a truncated string excerpt; `reverseTruncate: true` keeps the last `N` chars. + +Use `promptMaxChars` when partial data is worse than no data. Use `keepInPromptChars` when a prefix or suffix alone is useful. The two options are mutually exclusive on one field. + +### Auto-upgrade defaults + +`autoUpgrade` is ON by default: the agent applies both knobs above on the user's behalf based on character counts, so forgetting them no longer floods prompts. + +- Function discovery: when `functionDiscovery` is left unset and the estimated inline docs of discoverable functions exceed ~10k chars, discovery is enabled automatically. An explicit `functionDiscovery: true | false` always wins. +- Context fields: per run, an undeclared input value whose serialized size exceeds 8k chars is kept runtime-only like a declared context field — the prompt gets a 1,200-char truncated preview plus a `contextMetadata` entry, while the full value stays addressable as `inputs.` in the code runtime (the responder stage gets the same preview). Fields declared in `contextFields` keep their declared config. + +```typescript +autoUpgrade?: boolean | { + functionDiscovery?: boolean | { aboveFunctionDocChars?: number }; // default 10_000 + contextFields?: boolean | { promoteAboveChars?: number; previewChars?: number }; // 8_000 / 1_200 +} +``` + +Rules: + +- Set `autoUpgrade: false` (or disable one side) to restore fully manual behavior. +- Values in required non-string fields (arrays, objects, numbers, media) are never auto-promoted — declare those in `contextFields` explicitly when they can be large. +- Each promotion emits a `field_auto_promoted` context event (`onContextEvent`) with the field name, original size, and preview size; use it to observe what was kept out of the prompt. + +## Public Surface + +Use these method groups as the compact AxAgent surface map: + +- Running: `forward(ai, values, options?)` and `streamingForward(ai, values, options?)`. +- Forward-time agent options: `skills`, `onUsedMemories`, and `onUsedSkills`; use `ax-agent-memory-skills` for details. +- State and control: `getState()`, `setState(state?)`, `getContextMap()`, `setContextMap(map?)`, `stop()`, `getSignature()`, `setSignature(signature)`, `getFunction()`, `getId()`, and `setId(id)`. Context-map evolve policy lives on `AxAgentContextMap` (`infiniteEvolve`, `evolveSteps`, `maxChars`), not on the agent config. See [`src/examples/rlm-context-map-live.ts`](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/rlm-context-map-live.ts) for provider-backed persistence and finite-evolve usage. +- Observability: `getChatLog()`, `getUsage()`, `getStagedUsage()`, `resetUsage()`, and `getTraces()`; use `ax-agent-observability` for details. +- Demos and tuning: `setDemos(...)`, `namedPrograms()`, `namedProgramInstances()`, `optimize(...)`, `applyOptimization(...)`, `getOptimizableComponents()`, and `applyOptimizedComponents(...)`; use `ax-agent-optimize` for tuning details. + +Rules: + +- `getFunction()` requires `agentIdentity` because the agent needs function metadata when used as a child tool. +- Prefer `.forward(...)` for normal runs and `.streamingForward(...)` only when the caller needs streamed responder output. +- `setSignature(...)` must preserve configured `contextFields`; it throws if a configured context field is missing from the new signature. +- Treat low-level optimization component methods as advanced hooks; normal examples should use `agent.optimize(...)` and `agent.applyOptimization(...)`. + +## Tuning Hand-off + +When the user wants `agent.optimize(...)`, judge configuration, eval datasets, saved optimization artifacts, or optimization guidance, use `ax-agent-optimize`. + +Keep this skill focused on building and running agents. For tuning work: + +- use eval-safe tools +- treat `judgeOptions` as part of the optimize workflow +- choose an objective `metric` when scoring is mechanical; use the built-in judge only when run quality needs qualitative review +- keep runtime authoring guidance here and optimization guidance in `ax-agent-optimize` + +## Examples + +Fetch these for full working code: + +- [Agent](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/agent.ts) - basic agent +- [Functions](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/function.ts) - function validation +- [Food Search](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/food-search.ts) - API tools +- [Smart Home](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/smart-home.ts) - state management +- [Customer Support](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/customer-support.ts) - classification agent +- [Abort Patterns](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/abort-patterns.ts) - abort handling +- [Smart Defaults Agent](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/typescript/long-agents/smart-defaults-agent.ts) - auto-upgrade context promotion, relevance hints, and runtime tools + +RLM examples are listed in `ax-agent-rlm`. Memory/skills examples are listed in `ax-agent-memory-skills`. + +## Do Not Generate + +- Do not use `new AxAgent(...)` for new code unless explicitly required. +- Do not assume child agents are always under `agents.*`. +- Do not guess function names in discovery mode. +- Do not write a full multi-step RLM actor program in one turn; use `ax-agent-rlm`. +- Do not combine `console.log(...)` with `final(...)`. +- Do not add `bubbleErrors` for ordinary recoverable tool errors. +- Do not call `discover()` from the distiller or responder stages. +- Do not assign or inspect the return value of `await discover(...)`; read the next prompt instead. +- Do not loop `discover()` calls or wrap them in `Promise.all`. diff --git a/.claude/skills/ax-ai/SKILL.md b/.claude/skills/ax-ai/SKILL.md new file mode 100644 index 0000000..ffdb7cd --- /dev/null +++ b/.claude/skills/ax-ai/SKILL.md @@ -0,0 +1,382 @@ +--- +name: ax-ai +description: This skill helps an LLM generate correct AI provider setup and configuration code using @ax-llm/ax. Use when the user asks about ai(), providers, models, presets, embeddings, batch audio with ai.transcribe() or ai.speak(), extended thinking, context caching, or mentions OpenAI/Anthropic/Google/Azure/DeepSeek/Mistral/Cohere/Reka/Grok with @ax-llm/ax. +version: "23.0.0" +--- + +# AI Provider Codegen Rules (@ax-llm/ax) + +Use this skill to generate AI provider setup, configuration, and chat code. Prefer short, modern, copyable patterns. Do not write tutorial prose unless the user explicitly asks for explanation. + +## Quick Setup + +```typescript +import { ai } from '@ax-llm/ax'; + +const openai = ai({ name: 'openai', apiKey: 'sk-...' }); +const claude = ai({ name: 'anthropic', apiKey: 'sk-ant-...' }); +const gemini = ai({ name: 'google-gemini', apiKey: 'AIza...' }); +const azure = ai({ name: 'azure-openai', apiKey: 'your-key', resourceName: 'your-resource', deploymentName: 'gpt-5-4-mini' }); +const deepseek = ai({ name: 'deepseek', apiKey: 'sk-...' }); +const mistral = ai({ name: 'mistral', apiKey: 'your-key' }); +const cohere = ai({ name: 'cohere', apiKey: 'your-key' }); +const custom = ai({ + name: 'openai', + apiKey: process.env.PROVIDER_API_KEY, + apiURL: 'https://example.com/v1', + config: { model: 'provider/model-name' }, +}); +const reka = ai({ name: 'reka', apiKey: 'your-key' }); +const grok = ai({ name: 'grok', apiKey: 'your-key' }); +const compatible = ai({ name: 'openai', apiKey: 'key', apiURL: 'https://api.example.com/v1', config: { model: 'provider/model' } }); +``` + + +WebLLM is browser-only and requires a host-created WebLLM engine. The host +loads or reloads models with WebLLM APIs such as `CreateMLCEngine(...)`; Ax +only forwards chat requests to that loaded engine. Do not present WebLLM as a +portable AxIR provider or a server-side default. + +```typescript +import { ai, AxAIWebLLMModel } from '@ax-llm/ax'; + +const engine = await CreateMLCEngine(AxAIWebLLMModel.Llama32_3B_Instruct); +const llm = ai({ + name: 'webllm', + engine, + config: { + model: AxAIWebLLMModel.Llama32_3B_Instruct, + stream: false, + supportsFunctions: false, + }, +}); +``` + + +## Model Presets + +```typescript +import { ai, AxAIGoogleGeminiModel } from '@ax-llm/ax'; + +const gemini = ai({ + name: 'google-gemini', + apiKey: process.env.GOOGLE_APIKEY!, + config: { model: 'simple' }, + models: [ + { key: 'tiny', model: AxAIGoogleGeminiModel.Gemini31FlashLite, description: 'Fast + cheap', config: { maxTokens: 1024, temperature: 0.3 } }, + { key: 'simple', model: AxAIGoogleGeminiModel.Gemini35Flash, description: 'Balanced', config: { temperature: 0.6 } }, + ], +}); + +await gemini.chat({ model: 'tiny', chatPrompt: [{ role: 'user', content: 'Hi' }] }); +``` + +## Model Catalog + +```typescript +import { axGetSupportedAIModels } from '@ax-llm/ax'; + +const providers = axGetSupportedAIModels(); +const openai = providers.find((provider) => provider.name === 'openai'); +console.log(openai?.models[0]?.promptTokenCostPer1M); + +const textProviders = axGetSupportedAIModels({ type: 'text' }); +const embeddingProviders = axGetSupportedAIModels({ type: 'embeddings' }); +``` + +Use `axGetSupportedAIModels()` to build provider/model selectors before creating an `ai(...)` instance. It returns bundled static metadata: provider names, display names, default models, raw `AxModelInfo` pricing/details, model type (`'text'`, `'embeddings'`, `'code'`, or `'audio'`), and normalized capability flags for thinking, thoughts, structured outputs, audio, temperature, and top-p support. Provider groups and models are sorted cheapest to most expensive based on bundled input + output token pricing; unpriced models sort last. + +Filter with `{ type: 'all' | 'text' | 'embeddings' | 'code' | 'audio' }` or an array of those values. The `'text'` filter includes code-capable models; use `'code'` to show only code-first models. + +Dynamic providers such as Azure OpenAI deployments are marked with `isDynamic: true` and may have an empty or static-limited model list. + +## Chat + +```typescript +const res = await llm.chat({ + chatPrompt: [ + { role: 'system', content: 'You are concise.' }, + { role: 'user', content: 'Write a haiku about the ocean.' }, + ], +}); +console.log(res.results[0]?.content); +``` + +## Batch Audio + +Use `ai.transcribe(...)` for batch speech-to-text and `ai.speak(...)` for batch text-to-speech. These are separate from conversational `.chat()` audio config. + +```typescript +const transcript = await llm.transcribe({ + audio: { data: base64Wav, format: 'wav' }, + model: 'gpt-4o-mini-transcribe', + language: 'en', +}); + +const speech = await llm.speak({ + text: transcript.text, + model: 'gpt-4o-mini-tts', + voice: 'alloy', + format: 'mp3', +}); + +console.log(transcript.text); +console.log(speech.data); +``` + +Providers without the requested audio endpoint throw `AxMediaNotSupportedError`. Use `speech` forward options for signature audio artifacts and `modelConfig.audio` for conversational chat audio. + +## Common Options + +- `stream` (boolean): enable SSE; true by default +- `thinkingTokenBudget`: `'minimal'` | `'low'` | `'medium'` | `'high'` | `'highest'` | `'none'` +- `showThoughts`: include thoughts in output +- `functionCallMode`: `'auto'` | `'native'` | `'prompt'` +- `debug`, `logger`, `tracer`, `rateLimiter`, `timeout` + +## Global Runtime Defaults + +Use `axGlobals` when the app wants one live default for AI requests, generator runs, flows, or metrics: + +```typescript +import { ai, axGlobals, axCreateDefaultColorLogger } from '@ax-llm/ax'; +import { trace } from '@opentelemetry/api'; + +axGlobals.tracer = trace.getTracer('my-app'); +axGlobals.debug = true; +axGlobals.logger = axCreateDefaultColorLogger(); +axGlobals.customLabels = { service: 'api' }; + +const llm = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY! }); +``` + +Rules: + +- `axGlobals.tracer`, `meter`, `logger`, `debug`, `abortSignal`, and `customLabels` are live runtime defaults; future calls read the current value even if the AI instance already exists. +- Precedence is: per-call options, then explicit AI/service options, then current `axGlobals`, then built-in defaults. +- `customLabels` merge from globals to service to call options; later sources override earlier keys. +- `abortSignal` values are merged, so either a global shutdown signal or a local request signal can cancel the request. + +## DeepSeek Notes + +```typescript +import { ai, AxAIDeepSeekModel } from '@ax-llm/ax'; + +const deepseek = ai({ + name: 'deepseek', + apiKey: process.env.DEEPSEEK_APIKEY!, + config: { model: AxAIDeepSeekModel.DeepSeekV4Flash }, +}); +``` + +DeepSeek's current API models are `deepseek-v4-flash` and `deepseek-v4-pro`. +The deprecated `deepseek-chat` and `deepseek-reasoner` aliases are retained for +compatibility until DeepSeek removes them on 2026-07-24. + +DeepSeek V4 supports thinking mode. Ax sends `thinking: { type: "disabled" }` +by default to preserve non-thinking behavior, and enables it when +`thinkingTokenBudget` is set. Ax maps lower budget levels to DeepSeek's `high` +effort and maps `highest` to `max`. DeepSeek V4 thinking models support tools, +but reject the `tool_choice` request parameter, so Ax omits forced/auto tool +choice for `deepseek-v4-pro`, `deepseek-v4-flash`, and `deepseek-reasoner` +while still sending tool definitions. + +## Extended Thinking + +```typescript +import { ai, AxAIAnthropicModel } from '@ax-llm/ax'; + +const claude = ai({ + name: 'anthropic', + apiKey: process.env.ANTHROPIC_APIKEY!, + config: { model: AxAIAnthropicModel.Claude48Opus }, +}); + +const res = await claude.chat( + { chatPrompt: [{ role: 'user', content: 'Solve step by step...' }] }, + { thinkingTokenBudget: 'medium', showThoughts: true }, +); +console.log(res.results[0]?.thought); +console.log(res.results[0]?.content); +``` + +### Budget Levels + +| Level | Anthropic (tokens) | Gemini (tokens) | +|---|---|---| +| `'none'` | disabled | minimal | +| `'minimal'` | 1,024 | 200 | +| `'low'` | 5,000 | 800 | +| `'medium'` | 10,000 | 5,000 | +| `'high'` | 20,000 | 10,000 | +| `'highest'` | 32,000 | 24,500 | + +### Anthropic Model-Specific Behavior + +- Opus 4.8 and 4.7: adaptive thinking, effort levels including `'xhigh'`, + no manual `budget_tokens`, and no `temperature` / `topP` / `topK`. +- Opus 4.6: adaptive thinking, effort levels +- Opus 4.5: budget_tokens + effort levels (capped at `'high'`) +- Other thinking models: budget tokens only + +Anthropic `modelConfig.effort` can be set directly on a request. Fast mode and +task budgets are Anthropic-only opt-ins; `taskBudget.total` must be at least +20,000 tokens. + +```typescript +const res = await claude.chat({ + chatPrompt: [{ role: 'user', content: 'Review this migration plan.' }], + modelConfig: { + effort: 'xhigh', + speed: 'fast', + taskBudget: { type: 'tokens', total: 64_000 }, + }, +}); +``` + +### Custom Thinking Levels + +```typescript +const claude = ai({ + name: 'anthropic', + apiKey: '...', + config: { + model: AxAIAnthropicModel.Claude48Opus, + thinkingTokenBudgetLevels: { + minimal: 2048, + low: 8000, + medium: 16000, + high: 25000, + highest: 40000, + }, + effortLevelMapping: { + minimal: 'low', + low: 'medium', + medium: 'high', + high: 'high', + highest: 'max', + }, + }, +}); +``` + +## Embeddings + +```typescript +const { embeddings } = await llm.embed({ + texts: ['hello', 'world'], + embedModel: 'text-embedding-005', +}); +``` + +## Context Caching + +```typescript +const result = await gen.forward(llm, { code, language }, { + mem, + sessionId: 'code-review-session', + contextCache: { + ttlSeconds: 3600, + cacheBreakpoint: 'after-examples', + }, +}); +``` + +Breakpoint values: `'system'` | `'after-functions'` | `'after-examples'` + +Provider behavior: + +- Google Gemini: explicit caching with cache resource ID, auto TTL refresh +- Anthropic: implicit via `cache_control` markers + +### External Registry (serverless) + +```typescript +const registry: AxContextCacheRegistry = { + get: async (key) => { /* redis.get */ }, + set: async (key, entry) => { /* redis.set */ }, +}; +``` + +## AWS Bedrock + +```typescript +import { AxAIBedrock, AxAIBedrockModel } from '@ax-llm/ax-ai-aws-bedrock'; + +const bedrock = new AxAIBedrock({ + region: 'us-east-2', + fallbackRegions: ['us-west-2'], + config: { model: AxAIBedrockModel.ClaudeOpus45 }, +}); +``` + +## Vercel AI SDK Integration + +```typescript +import { generateText } from 'ai'; +import { ai } from '@ax-llm/ax'; +import { AxAIProvider } from '@ax-llm/ax-ai-sdk-provider'; + +const axAI = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY ?? '', +}); +const model = new AxAIProvider(axAI); + +const result = await generateText({ + model, + prompt: 'Hello!', +}); +``` + +## MCP + AxJSRuntime + +```typescript +import { AxMCPClient } from '@ax-llm/ax'; +import { axCreateMCPStdioTransport } from '@ax-llm/ax-tools'; + +const transport = axCreateMCPStdioTransport({ + command: 'npx', + args: ['-y', '@anthropic/mcp-server-filesystem'], +}); +const client = new AxMCPClient(transport); +``` + +## Critical Rules + +- Use `ai()` factory for all providers. +- Provider names: `'openai'`, `'openai-responses'`, `'anthropic'`, `'google-gemini'`, `'azure-openai'`, `'mistral'`, `'cohere'`, `'deepseek'`, `'reka'`, `'grok'` +- Thinking constraints on Anthropic: Opus 4.8/4.7 omit `temperature`, `topP`, + and `topK`; older thinking models ignore `temperature` and `topK`, with + `topP` only sent if >= 0.95. +- Bedrock uses `new AxAIBedrock()`, not `ai()`. +- Vercel AI SDK uses `AxAIProvider` wrapper. + +## Examples + +Fetch these for full working code: + +- [Embeddings](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/embed.ts) — embedding generation +- [Anthropic Thinking](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/anthropic-thinking-function.ts) — extended thinking with functions +- [Anthropic Thinking Separation](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/anthropic-thinking-separation.ts) — thinking separation +- [Anthropic Web Search](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/anthropic-web-search.ts) — Anthropic web search +- [OpenAI Web Search](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/openai-web-search.ts) — OpenAI web search +- [OpenAI Responses](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/openai-responses.ts) — OpenAI responses API +- [o3 Reasoning](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/reasoning-o3-example.ts) — o3 reasoning +- [Gemini Context Cache](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/gemini-context-cache.ts) — Gemini context caching +- [Gemini Files](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/gemini-file-support.ts) — Gemini file handling +- [Grok Live Search](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/grok-live-search.ts) — Grok live search +- [OpenAI-Compatible](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/openai-compatible.ts) — custom OpenAI-compatible base URL +- [Vertex AI Auth](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/vertex-auth-example.ts) — Vertex AI authentication +- [MCP Stdio](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/mcp-client-memory.ts) — MCP stdio transport +- [MCP HTTP](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/mcp-client-pipedream.ts) — MCP HTTP transport +- [Telemetry](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/telemetry.ts) — OpenTelemetry tracing +- [Multi-Modal](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/multi-modal.ts) — image handling + +## Do Not Generate + +- Do not use `new AxAIOpenAI(...)` or similar class constructors for standard providers; use `ai()`. +- Do not hardcode provider class names when `ai({ name: ... })` covers the provider. +- Do not mix `thinkingTokenBudget` with explicit `temperature` on Anthropic thinking models. +- Do not use `ai()` for AWS Bedrock; use `new AxAIBedrock()`. +- Do not omit `resourceName` and `deploymentName` for Azure OpenAI. diff --git a/.claude/skills/ax-audio/SKILL.md b/.claude/skills/ax-audio/SKILL.md new file mode 100644 index 0000000..7063ae5 --- /dev/null +++ b/.claude/skills/ax-audio/SKILL.md @@ -0,0 +1,373 @@ +--- +name: ax-audio +description: This skill helps an LLM generate correct audio code with @ax-llm/ax. Use when the user asks about ai.transcribe(), ai.speak(), signature audio inputs or outputs, agent audio behavior, .chat() conversational audio, OpenAI audio or realtime models, Gemini Live native audio, Grok Voice Agent models, voices, formats, transcripts, or how audio fits with structured outputs. +version: "23.0.0" +--- + +# Audio I/O Codegen Rules (@ax-llm/ax) + +Use this skill for audio in Ax. Pick the smallest audio surface that matches the job: + +- Use `ai.transcribe(...)` for batch speech-to-text. +- Use `ai.speak(...)` for batch text-to-speech. +- Use `speech:audio` signature outputs for structured programs that should return synthesized audio artifacts. +- Use `.chat()` audio config for conversational or realtime audio turns. + +## Core Rules + +- Input `:audio` is an audio input value: `{ data, format?, mimeType?, sampleRate?, channels? }`. +- Output `:audio` is a scripted audio artifact. The model returns plain text for that field; Ax synthesizes it after structured output parsing. +- Output audio JSON schema is model-facing `string`, not a binary object. +- Agents transcribe input audio fields before planner/executor/responder stages by default, so agent stages see text instead of base64 audio. +- Realtime and conversational audio still use `.chat()` and `modelConfig.audio`. +- Batch signature audio artifacts use forward-time `speech` options, not `modelConfig.audio`. + +## Direct Batch APIs + +```typescript +import { ai } from '@ax-llm/ax'; + +const llm = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY! }); + +const transcript = await llm.transcribe({ + audio: { data: base64Wav, format: 'wav' }, + model: 'gpt-4o-mini-transcribe', + language: 'en', + prompt: 'Product support call', +}); + +const speech = await llm.speak({ + text: transcript.text, + model: 'gpt-4o-mini-tts', + voice: 'alloy', + format: 'mp3', +}); + +console.log(transcript.text); +console.log(speech.data); +console.log(speech.transcript); +``` + +Providers without the requested batch audio capability throw `AxMediaNotSupportedError`. + +## Signature Audio Artifacts + +```typescript +import { ai, ax } from '@ax-llm/ax'; + +const llm = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY! }); +const say = ax('question:string -> speech:audio, summary:string'); + +const result = await say.forward( + llm, + { question: 'Explain retries in one sentence.' }, + { + speech: { + speak: { voice: 'alloy', format: 'mp3' }, + fields: { + speech: { voice: 'alloy' }, + }, + }, + } +); + +console.log(result.summary); +console.log(result.speech.data); +console.log(result.speech.mimeType); +console.log(result.speech.transcript); +``` + +The model emits a text script for `speech`; Ax replaces it with `AxChatAudioOutput` after result selection. If the field already contains an audio artifact with `{ data }` or `{ id }`, Ax leaves it alone. + +## Agent Audio Inputs + +```typescript +import { agent, ai } from '@ax-llm/ax'; + +const llm = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY! }); + +const voiceAgent = agent( + 'recording:audio, question:string -> speech:audio, summary:string', + { + agentIdentity: { + name: 'Voice Assistant', + description: 'Answers spoken requests with spoken and written output', + }, + contextFields: [], + } +); + +const result = await voiceAgent.forward( + llm, + { + recording: { data: base64Wav, format: 'wav' }, + question: 'What should I do next?', + }, + { + speech: { + transcribe: { model: 'gpt-4o-mini-transcribe' }, + speak: { voice: 'alloy', format: 'mp3' }, + }, + } +); + +console.log(result.summary); +console.log(result.speech.data); +``` + +The agent runtime transcribes `recording` first and passes the transcript through the internal agent stages. Use direct `ax(...)` or `.chat()` when you specifically want native audio understanding in the model call. + +## Conversational `.chat()` Audio + +Use `modelConfig.audio` for conversational audio turns where audio is part of the chat response instead of a structured signature field. + +```typescript +const res = await llm.chat({ + chatPrompt: [{ role: 'user', content: 'Say hello out loud.' }], + modelConfig: { + audio: { output: { enabled: true, voice: 'alloy', format: 'wav' } }, + }, +}); + +console.log(res.results[0]?.content); +console.log(res.results[0]?.audio?.data); +console.log(res.results[0]?.audio?.transcript); +``` + +## Config Shape + +```typescript +type AxAudioFormat = + | 'wav' + | 'mp3' + | 'flac' + | 'opus' + | 'aac' + | 'pcm16' + | 'pcm' + | 'ogg' + | 'raw' + | 'mulaw' + | 'ulaw' + | 'alaw'; + +type AxSpeechConfig = { + transcribe?: { + model?: string; + language?: string; + prompt?: string; + }; + speak?: { + model?: string; + voice?: string; + format?: AxAudioFormat; + }; + fields?: Record< + string, + { + model?: string; + voice?: string; + format?: AxAudioFormat; + } + >; +}; +``` + +## OpenAI Defaults + +Use `axAIOpenAIAudioDefaultConfig()` for OpenAI request-based audio chat: + +- model: `gpt-audio-mini` +- output enabled +- voice: `alloy` +- output format: `wav` +- transcript enabled +- streaming disabled by default +- audio input formats: `wav`, `mp3` +- audio output formats: `wav`, `mp3`, `flac`, `opus`, `aac`, `pcm16` + +```typescript +import { ai, axAIOpenAIAudioDefaultConfig } from '@ax-llm/ax'; + +const openai = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, + config: axAIOpenAIAudioDefaultConfig(), +}); + +const res = await openai.chat({ + chatPrompt: [ + { + role: 'user', + content: [ + { type: 'text', text: 'What is in this recording?' }, + { type: 'audio', data: base64Wav, format: 'wav' }, + ], + }, + ], +}); + +console.log(res.results[0]?.content); +console.log(res.results[0]?.audio?.data); +``` + +Use `axAIOpenAIRealtimeDefaultConfig()` for OpenAI realtime speech-to-speech: + +- model: `gpt-realtime-2` +- output enabled +- voice: `marin` +- output format: `pcm16` +- input default: `audio/pcm`, mono, 24000 Hz +- turn timeout: `30000` +- streaming disabled by default + +Use `axAIOpenAIRealtimeTranscriptionDefaultConfig()` for realtime transcript deltas: + +- model: `gpt-realtime-whisper` +- input default: `audio/pcm`, mono, 24000 Hz +- output audio disabled; transcript text is returned on `content` + +Realtime models use a one-turn WebSocket call under `.chat()`. In Node, pass a WebSocket constructor through request options: + +```typescript +import WebSocket from 'ws'; +import { ai, axAIOpenAIRealtimeDefaultConfig } from '@ax-llm/ax'; + +const openai = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, + config: axAIOpenAIRealtimeDefaultConfig(), +}); + +const stream = await openai.chat( + { + chatPrompt: [{ role: 'user', content: 'Say hello out loud.' }], + }, + { stream: true, webSocket: WebSocket } +); +``` + +For follow-up turns, keep the assistant audio reference in history: + +```typescript +await openai.chat({ + chatPrompt: [ + { role: 'assistant', audio: { id: previousAudioId } }, + { role: 'user', content: 'Repeat that more slowly.' }, + ], +}); +``` + +## Gemini Live Defaults + +Use `axAIGoogleGeminiLiveAudioDefaultConfig()` for Gemini native audio: + +- model: `gemini-2.5-flash-native-audio-preview-12-2025` +- output enabled +- voice: `Kore` +- output format: `pcm16` +- output sample rate: `24000` +- input default: `audio/pcm;rate=16000`, mono +- transcript enabled +- turn timeout: `30000` +- streaming disabled by default + +```typescript +import { ai, axAIGoogleGeminiLiveAudioDefaultConfig } from '@ax-llm/ax'; + +const gemini = ai({ + name: 'google-gemini', + apiKey: process.env.GOOGLE_APIKEY!, + config: axAIGoogleGeminiLiveAudioDefaultConfig(), +}); + +const res = await gemini.chat({ + chatPrompt: [ + { + role: 'user', + content: [ + { type: 'text', text: 'Answer this spoken question.' }, + { + type: 'audio', + data: base64Pcm16, + format: 'pcm16', + sampleRate: 16000, + channels: 1, + }, + ], + }, + ], +}); + +console.log(res.results[0]?.content); +console.log(res.results[0]?.audio?.data); +``` + +Gemini Live uses a one-turn WebSocket call under `.chat()`. It expects PCM input for native audio turns; use `format: 'pcm16'` or `mimeType: 'audio/pcm;rate=16000'`. + +## Grok Voice Defaults + +Use `axAIGrokVoiceDefaultConfig()` for xAI Grok Voice Agent: + +- model: `grok-voice-think-fast-1.0` +- output enabled +- voice: `eve` +- output format: `pcm16` +- output sample rate: `24000` +- input default: `audio/pcm`, mono, 24000 Hz +- transcript enabled +- turn timeout: `30000` +- streaming disabled by default + +```typescript +import WebSocket from 'ws'; +import { ai, axAIGrokVoiceDefaultConfig } from '@ax-llm/ax'; + +const grok = ai({ + name: 'grok', + apiKey: process.env.GROK_API_KEY!, + config: axAIGrokVoiceDefaultConfig(), +}); + +const res = await grok.chat( + { + chatPrompt: [{ role: 'user', content: 'Say hello out loud.' }], + }, + { webSocket: WebSocket } +); + +console.log(res.results[0]?.content); +console.log(res.results[0]?.audio?.data); +``` + +Grok Voice uses a one-turn WebSocket call under `.chat()`. It expects PCM input for spoken input turns; use `format: 'pcm16'` or `mimeType: 'audio/pcm'`. + +## Streaming Audio + +OpenAI audio chat, OpenAI Realtime, Gemini Live, and Grok Voice all default to non-streaming, but each can stream deltas when you pass `{ stream: true }`. + +```typescript +const stream = await llm.chat( + { + chatPrompt: [{ role: 'user', content: 'Say hello.' }], + }, + { stream: true } +); + +for await (const chunk of stream) { + const audio = chunk.results[0]?.audio; + if (audio?.isDelta) { + playAudioChunk(audio.data); + } +} +``` + +## Structured Outputs + +Use signature audio outputs for structured speech artifacts: + +```typescript +const gen = ax('question:string -> answer:string, speech:audio'); +``` + +Use `.chat()` audio when the response itself is a conversational audio turn. Do not combine `.chat()` audio output with provider-native structured response formats unless that provider explicitly supports the combination. diff --git a/.claude/skills/ax-flow/SKILL.md b/.claude/skills/ax-flow/SKILL.md new file mode 100644 index 0000000..6121dbf --- /dev/null +++ b/.claude/skills/ax-flow/SKILL.md @@ -0,0 +1,442 @@ +--- +name: ax-flow +description: This skill helps an LLM generate correct AxFlow workflow code using @ax-llm/ax. Use when the user asks about flow(), AxFlow, workflow orchestration, parallel execution, DAG workflows, conditional routing, map/reduce patterns, or multi-node AI pipelines. +version: "23.0.0" +--- + +# AxFlow Codegen Rules (@ax-llm/ax) + +Use this skill to generate `AxFlow` workflow code. Prefer short, modern, copyable patterns. Do not write tutorial prose unless the user explicitly asks for explanation. + +## Use These Defaults + +- Use `flow()` factory, not `new AxFlow()`. +- Import: `import { ai, flow, f } from '@ax-llm/ax';` +- `autoParallel: true` is the default; independent executes and derives run in parallel when their metadata reads/writes are known and non-conflicting. +- Node results are stored as `${nodeName}Result` in state. +- Always define `.node()` before `.execute()` for that node. +- Use `.returns()` (or `.r()`) as the last step to lock the output type. +- Use descriptive node names: `documentSummarizer`, not `proc1`. +- Use descriptive field names: `userInput`, `responseText`, not `text`, `result`. + +## Critical Rules + +- Use `flow()` factory syntax for new code. +- Node results in state follow the pattern `state.${nodeName}Result.${fieldName}`. +- `.execute()` maps current state to node inputs; `.map()` transforms state without AI calls. +- `.returns()` maps final state to the flow output type. +- Always define nodes before executing them; reversed order throws at runtime. +- Keep state flat; avoid deep nesting in `.map()`. +- Ensure loop conditions can change to avoid infinite loops. +- Structure independent executes to maximize safe auto-parallelization. +- Use `flow()` for typed flows. +- Aliases: `.n()` = `.node()`, `.nx()` = `.nodeExtended()`, `.m()` = `.map()`, `.r()` = `.returns()`. + +## Canonical Pattern + +```typescript +import { ai, flow } from '@ax-llm/ax'; + +const llm = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY! }); + +const wf = flow<{ userInput: string }, { responseText: string }>() + .node('testNode', 'userInput:string -> responseText:string') + .execute('testNode', (state) => ({ userInput: state.userInput })) + .returns((state) => ({ responseText: state.testNodeResult.responseText })); + +const result = await wf.forward(llm, { userInput: 'Hello world' }); +console.log(result.responseText); +``` + +## Factory Options + +```typescript +// Basic +const wf = flow(); + +// With options +const wf = flow({ autoParallel: false }); + +// Typed +const wf = flow(); + +// Typed with options +const wf = flow({ autoParallel: true, batchSize: 5 }); +``` + +## State Evolution + +State grows with each executed node. Results are stored as `${nodeName}Result`: + +```typescript +// Initial state: { userInput: 'Hello' } +flow.execute('processor', (state) => ({ input: state.userInput })); +// State: { userInput: 'Hello', processorResult: { output: '...' } } + +flow.execute('analyzer', (state) => ({ text: state.processorResult.output })); +// State: { ..., analyzerResult: { sentiment: '...', confidence: 0.8 } } +``` + +## Node Definition + +```typescript +// String signature (creates AxGen automatically) +flow.node('processor', 'input:string -> output:string'); + +// Multiple outputs +flow.node('analyzer', 'text:string -> sentiment:string, confidence:number'); + +// Array outputs +flow.node('extractor', 'documentText:string -> entities:string[]'); + +// Short alias +flow.n('processor', 'input:string -> output:string'); +``` + +## Extended Nodes (nx) + +Add fields to a base signature without rewriting it: + +```typescript +import { f, flow } from '@ax-llm/ax'; + +// Chain-of-thought reasoning +flow.nx('reasoner', 'question:string -> answer:string', { + prependOutputs: [ + { name: 'reasoning', type: f.internal(f.string('Step-by-step reasoning')) }, + ], +}); + +// Add confidence scoring +flow.nx('analyzer', 'input:string -> result:string', { + appendOutputs: [{ name: 'confidence', type: f.number('Confidence 0-1') }], +}); + +// Add optional context input +flow.nx('processor', 'query:string -> response:string', { + appendInputs: [{ name: 'context', type: f.optional(f.string('Extra context')) }], +}); +``` + +Extension options: `prependInputs`, `appendInputs`, `prependOutputs`, `appendOutputs`. + +## Execute With Input Mapping + +```typescript +flow.execute('summarizer', (state) => ({ documentText: state.document })); + +// With AI override (use a different model for this node) +flow.execute('processor', (state) => ({ input: state.data }), { ai: alternativeAI }); +``` + +## Map (State Transformation) + +Use `map()` for data shaping without AI calls: + +```typescript +// Sync +flow.map((state) => ({ ...state, upperText: state.rawText.toUpperCase() })); + +// Async +flow.map(async (state) => { + const data = await fetchFromAPI(state.query); + return { ...state, enrichedData: data }; +}); + +// Parallel async transforms +flow.map([ + async (state) => ({ ...state, result1: await api1(state.data) }), + async (state) => ({ ...state, result2: await api2(state.data) }), +], { parallel: true }); +``` + +## Returns (Final Output) + +```typescript +const wf = flow<{ input: string }>() + .map((state) => ({ ...state, upper: state.input.toUpperCase(), len: state.input.length })) + .returns((state) => ({ upper: state.upper, isLong: state.len > 20 })); + +// Result is typed as { upper: string; isLong: boolean } +const result = await wf.forward(llm, { input: 'test' }); +``` + +## Sequential Processing + +```typescript +const wf = flow<{ input: string }, { finalResult: string }>() + .node('step1', 'input:string -> intermediate:string') + .node('step2', 'intermediate:string -> output:string') + .execute('step1', (state) => ({ input: state.input })) + .execute('step2', (state) => ({ intermediate: state.step1Result.intermediate })) + .returns((state) => ({ finalResult: state.step2Result.output })); +``` + +## Auto-Parallel Execution + +Independent execute steps run in parallel automatically (`autoParallel: true` by default) when their metadata reads/writes are known and non-conflicting: + +```typescript +const wf = flow<{ text: string }, { combined: string }>() + .node('sentimentAnalyzer', 'text:string -> sentiment:string') + .node('topicExtractor', 'text:string -> topics:string[]') + .node('entityRecognizer', 'text:string -> entities:string[]') + // These three run in parallel (all depend only on state.text) + .execute('sentimentAnalyzer', (state) => ({ text: state.text })) + .execute('topicExtractor', (state) => ({ text: state.text })) + .execute('entityRecognizer', (state) => ({ text: state.text })) + // This waits for all three + .returns((state) => ({ + combined: JSON.stringify({ + sentiment: state.sentimentAnalyzerResult.sentiment, + topics: state.topicExtractorResult.topics, + entities: state.entityRecognizerResult.entities, + }), + })); + +// Inspect execution plan +const plan = wf.getExecutionPlan(); +console.log(plan.parallelGroups, plan.maxParallelism); +``` + +Planner rules: +- Independent `.execute()` and `.derive()` steps may parallelize. +- `.map()`, `.returns()`, `.branch()`, `.while()`, `.feedback()`, and explicit `.parallel()` are barriers. +- Branch, while, and feedback bodies still use the same planner internally. +- Use `autoParallel: false` when you need strict sequential execution. + +Disable auto-parallel: + +```typescript +const wf = flow({ autoParallel: false }); +// or per execution: +await wf.forward(llm, input, { autoParallel: false }); +``` + +## Conditional Branching + +```typescript +const wf = flow<{ query: string; expertMode: boolean }, { response: string }>() + .node('simple', 'query:string -> response:string') + .node('expert', 'query:string -> response:string') + .branch((state) => state.expertMode) + .when(true) + .execute('expert', (state) => ({ query: state.query })) + .when(false) + .execute('simple', (state) => ({ query: state.query })) + .merge() + .returns((state) => ({ + response: state.expertResult?.response ?? state.simpleResult?.response, + })); +``` + +After `.merge()`, only the taken branch's result exists; use optional chaining (`?.`) on untaken branch results. + +## While Loops + +```typescript +const wf = flow<{ content: string }, { finalContent: string }>() + .node('processor', 'content:string -> processedContent:string') + .node('qualityChecker', 'content:string -> qualityScore:number') + .map((state) => ({ currentContent: state.content, iteration: 0, qualityScore: 0 })) + .while((state) => state.iteration < 3 && state.qualityScore < 0.8) + .map((state) => ({ ...state, iteration: state.iteration + 1 })) + .execute('processor', (state) => ({ content: state.currentContent })) + .execute('qualityChecker', (state) => ({ + content: state.processorResult.processedContent, + })) + .map((state) => ({ + ...state, + currentContent: state.processorResult.processedContent, + qualityScore: state.qualityCheckerResult.qualityScore, + })) + .endWhile() + .returns((state) => ({ finalContent: state.currentContent })); +``` + +Rules: +- Every `.while()` needs a matching `.endWhile()`. +- Ensure the loop condition can change to avoid infinite loops. + +## Feedback Loops (label/feedback) + +```typescript +const wf = flow<{ prompt: string }, { result: string }>() + .node('gen', 'prompt:string -> result:string, quality:number') + .map((state) => ({ ...state, tries: 0 })) + .label('retry') + .map((state) => ({ ...state, tries: state.tries + 1 })) + .execute('gen', (state) => ({ prompt: state.prompt })) + .feedback((state) => state.genResult.quality < 0.9 && state.tries < 3, 'retry') + .returns((state) => ({ result: state.genResult.result })); +``` + +Rules: +- Define the label before referencing it in `.feedback()`. +- Always include a max-iteration guard to avoid infinite loops. + +## Explicit Parallel Sub-Flows + +```typescript +flow + .parallel([ + (sub) => sub.execute('analyzer1', (state) => ({ text: state.input })), + (sub) => sub.execute('analyzer2', (state) => ({ text: state.input })), + (sub) => sub.execute('analyzer3', (state) => ({ text: state.input })), + ]) + .merge('combinedResults', (r1, r2, r3) => ({ + a1: r1.analyzer1Result.analysis, + a2: r2.analyzer2Result.analysis, + a3: r3.analyzer3Result.analysis, + })); +``` + +## Derive (Batch/Array Processing) + +```typescript +const wf = flow<{ items: string[] }, { processed: string[] }>({ batchSize: 3 }) + .derive('processed', 'items', (item, index) => `processed-${item}-${index}`, { + batchSize: 2, + }); +``` + +## Dynamic AI Context (Multi-Model) + +Route nodes to different AI providers: + +```typescript +const fast = ai({ name: 'openai', apiKey: '...', config: { model: 'gpt-5.4-mini' } }); +const smart = ai({ name: 'anthropic', apiKey: '...' }); + +const wf = flow<{ text: string }, { out: string }>() + .node('draft', 'text:string -> out:string') + .node('refine', 'text:string -> out:string') + .execute('draft', (state) => ({ text: state.text }), { ai: fast }) + .execute('refine', (state) => ({ text: state.draftResult.out }), { ai: smart }) + .returns((state) => ({ out: state.refineResult.out })); +``` + +## Description and toFunction + +```typescript +const wf = flow<{ userQuestion: string }, { responseText: string }>() + .node('qa', 'userQuestion:string -> responseText:string') + .execute('qa', (state) => ({ userQuestion: state.userQuestion })) + .returns((state) => ({ responseText: state.qaResult.responseText })) + .description('Question Answerer', 'Answers user questions concisely.'); + +const fn = wf.toFunction(); +// fn.name, fn.parameters (JSON Schema), fn.func +``` + +## Instrumentation (Tracing) + +```typescript +import { ai, flow } from '@ax-llm/ax'; +import { context, trace } from '@opentelemetry/api'; + +const tracer = trace.getTracer('axflow'); +const llm = ai({ name: 'openai', apiKey: '...' }); + +const wf = flow<{ userQuestion: string }>() + .node('summarizer', 'documentText:string -> summaryText:string') + .execute('summarizer', (s) => ({ documentText: s.userQuestion })) + .returns((s) => ({ answer: s.summarizerResult.summaryText })); + +const result = await wf.forward(llm, { userQuestion: 'hi' }, { + tracer, + traceContext: context.active(), +}); +``` + +Flow tracing also respects live app-wide defaults: + +```typescript +import { axGlobals } from '@ax-llm/ax'; +import { metrics } from '@opentelemetry/api'; + +axGlobals.tracer = tracer; +axGlobals.meter = metrics.getMeter('axflow'); + +const result = await wf.forward(llm, { userQuestion: 'hi' }); +``` + +Rules: + +- `wf.forward(..., { tracer, meter })` overrides flow defaults and `axGlobals`. +- Constructor/factory flow defaults override `axGlobals`. +- If no local tracer or meter is provided, `AxFlow` reads current `axGlobals.tracer` and `axGlobals.meter`, creates a parent flow span, and propagates tracer/meter plus trace context to node forwards. +- `axGlobals.abortSignal` is merged with flow-level abort signals. + +## Program IDs and Demos + +```typescript +const wf = flow<{ input: string }>() + .node('summarizer', 'text:string -> summary:string') + .node('classifier', 'text:string -> category:string'); + +// Discover program IDs +console.log(wf.namedPrograms()); +// [{ id: 'root.summarizer', ... }, { id: 'root.classifier', ... }] + +// Set demos (TypeScript catches typos) +wf.setDemos([{ programId: 'root.summarizer', traces: [] }]); + +// Apply optimization +wf.applyOptimization(optimizedProgram); +``` + +For tuning a flow, use top-level `optimize(wf, train, metric, options)` from the +`ax-gepa` skill. There is no separate `flow.optimize(...)` helper. + +## Chat Logs + +`AxFlow.getChatLog()` returns a flat `readonly AxChatLogEntry[]` after `forward()`. Each child-node entry is tagged with `entry.name` so callers can filter by node: + +```typescript +const log = wf.getChatLog(); +for (const entry of log) { + console.log(entry.name, entry.model); +} +``` + +## Error Handling + +```typescript +try { + const result = await wf.forward(llm, input); +} catch (error) { + console.error('Flow execution failed:', error); +} +``` + +Common errors: +- `"Node 'x' not found"` -- define `.node()` before `.execute()`. +- `"endWhile() without matching while()"` -- every `.while()` needs `.endWhile()`. +- `"when() without matching branch()"` -- `.when()` must be inside `.branch()`/`.merge()`. +- `"merge() without matching branch()"` -- every `.branch()` needs `.merge()`. +- `"Label 'x' not found"` -- define `.label()` before `.feedback()` references it. + +## Examples + +Fetch these for full working code: + +- [Flow](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/ax-flow.ts) — complete flow usage +- [Auto-Parallel](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/ax-flow-auto-parallel.ts) — auto-parallelization +- [Async Map](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/ax-flow-async-map.ts) — async map transforms +- [Enhanced Demo](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/ax-flow-enhanced-demo.ts) — instance-based nodes +- [Flow as Function](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/ax-flow-to-function.ts) — flow as callable function +- [Fluent Builder](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/fluent-flow-example.ts) — fluent builder pattern +- [Load Balancing](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/balancer.ts) — load balancing + +## Do Not Generate + +- Do not use `new AxFlow(...)` for new code. +- Do not execute a node before defining it with `.node()`. +- Do not use removed terminal shapers like `.mapOutput()` or `.mo()`. +- Do not rely on broad signature inference from arbitrary transform source. Use explicit input/output generics and `.returns()` for the final output contract. +- Do not use generic field names like `text`, `result`, `data`, `input`, `output`. +- Do not create deep-nested state objects in `.map()`. +- Do not create loop conditions that can never change. +- Do not add unnecessary dependencies between executes (kills auto-parallelism). +- Do not forget to use optional chaining on branch results after `.merge()`. diff --git a/.claude/skills/ax-gen/SKILL.md b/.claude/skills/ax-gen/SKILL.md new file mode 100644 index 0000000..effd9a9 --- /dev/null +++ b/.claude/skills/ax-gen/SKILL.md @@ -0,0 +1,495 @@ +--- +name: ax-gen +description: This skill helps an LLM generate correct AxGen code using @ax-llm/ax. Use when the user asks about ax(), AxGen, generators, forward(), streamingForward(), validation, assertions, streaming assertions, field processors, step hooks, self-tuning, or structured outputs. +version: "23.0.0" +--- + +# AxGen Codegen Rules (@ax-llm/ax) + +Use this skill to generate `AxGen` code. Prefer short, modern, copyable patterns. Do not write tutorial prose unless the user explicitly asks for explanation. + +## Use These Defaults + +- Use `ax(...)` factory, not `new AxGen(...)`. +- Always pass an AI instance from `ai(...)` as the first argument to `forward()`. +- Streaming uses `streamingForward()`, not `forward()` with a stream option. +- Use schema validation for field shape and constraints. +- Use `addAssert(...)` for whole-output hard invariants with correction retries. +- Use `addStreamingAssert(...)` for partial streaming hard invariants with fail-fast per-attempt correction retries. +- Use `bestOfN(...)` / `refine(...)` for reward-scored complete outputs. +- Step hook mutations are applied at the next step boundary (pending pattern). +- `stopFunction` accepts a string or string[] for multiple stop functions. +- Multi-step continues until: all outputs filled, stop function called, or `maxSteps` reached. + +## Canonical Pattern + +```typescript +import { ai, ax, s } from '@ax-llm/ax'; + +const llm = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, +}); + +// Inline signature +const gen = ax('input:string -> output:string, reasoning:string'); + +// Reusable signature +const sig = s('question:string, context:string[] -> answer:string'); +const gen2 = ax(sig); + +// With options +const gen3 = ax('input -> output', { + description: 'A helpful assistant', + maxRetries: 3, + maxSteps: 10, + temperature: 0.7, +}); + +const result = await gen.forward(llm, { input: 'Hello world' }); +console.log(result.output); +``` + +### Signatures from zod / valibot / arktype + +`ax()` accepts any signature built with `f()`, and `f().input()` / `.output()` accept [Standard Schema v1](https://standardschema.dev) validators directly — per-field or a whole `z.object({...})`: + +```typescript +import { z } from 'zod'; +import { ax, f } from '@ax-llm/ax'; + +const gen = ax( + f() + .input(z.object({ + productName: z.string(), + buyerProfile: z.string(), + })) + .output(z.object({ + headline: z.string(), + recommendation: z.enum(['buy', 'wait', 'skip']), + })) + .build() +); +``` + +Constraints (`.min()`, `.email()`, `.regex()`) and custom logic (`.refine()`, `.transform()`, `.superRefine()`) execute in the normal validation/retry pipeline — at parse time on complete field values, including at field boundaries during streaming. For cache/internal hints pass companion options: `.input('ctx', z.string(), { cache: true })` or `.output('reasoning', z.string(), { internal: true })`. + +Define tool functions with zod the same way — `fn().arg()` / `.returns()` accept per-argument or whole-object schemas and infer the handler's argument type: + +```typescript +import { z } from 'zod'; +import { ax, fn } from '@ax-llm/ax'; + +const lookupProduct = fn('lookupProduct') + .description('Look up a product by name') + .arg(z.object({ + productName: z.string().min(1), + includeSpecs: z.boolean().optional(), + })) + .returns(z.object({ + price: z.number(), + inStock: z.boolean(), + rating: z.number().min(1).max(5), + })) + .handler(async ({ productName, includeSpecs }) => ({ + price: 79.99, + inStock: true, + rating: 4.3, + })) + .build(); + +const result = await gen.forward(llm, { ... }, { functions: [lookupProduct] }); +``` + +## Running AxGen + +### `forward()` + +```typescript +const result = await gen.forward(llm, { input: '...' }); + +// With options +const result = await gen.forward(llm, { input: '...' }, { + maxRetries: 5, + model: 'gpt-5.4-mini', + modelConfig: { temperature: 0.9, maxTokens: 1000 }, + debug: true, +}); +``` + +### Live Global Defaults + +`AxGen` respects `axGlobals` for app-wide runtime defaults: + +```typescript +import { axGlobals } from '@ax-llm/ax'; +import { trace } from '@opentelemetry/api'; + +const responseCache = new Map(); + +axGlobals.tracer = trace.getTracer('my-app'); +axGlobals.debug = true; +axGlobals.cachingFunction = async (key, value?) => { + if (value !== undefined) { + responseCache.set(key, value); + return; + } + return responseCache.get(key); +}; +``` + +Rules: + +- Tracing/logging precedence is: forward options, then generator options, then AI service options, then current `axGlobals`, then built-in defaults. +- `abortSignal` from `axGlobals` is merged with local forward signals. +- `customLabels` merge from globals to AI service to forward options. +- `cachingFunction` and `functionResultFormatter` also fall back to current `axGlobals` when local options do not provide them. + +### `streamingForward()` + +```typescript +const stream = gen.streamingForward(llm, { input: 'Write a long story' }); +for await (const chunk of stream) { + if (chunk.delta.output) process.stdout.write(chunk.delta.output); +} +``` + +## Stopping And Cancellation + +```typescript +import { AxAIServiceAbortedError } from '@ax-llm/ax'; + +const timer = setTimeout(() => gen.stop(), 3_000); + +try { + const result = await gen.forward(llm, { topic: 'Long document' }, { + abortSignal: AbortSignal.timeout(10_000), + }); +} catch (err) { + if (err instanceof AxAIServiceAbortedError) console.log('Aborted'); +} +``` + +Rules: + +- `gen.stop()` gracefully stops multi-step execution at the next step boundary. +- `abortSignal` cancels the underlying AI service call immediately. +- Catch `AxAIServiceAbortedError` when using either mechanism. + +## Validation, Selection, And Guards + +```typescript +import { ax, bestOfN, f } from '@ax-llm/ax'; +import { z } from 'zod'; + +// Schema validation: output shape and field validity. +const gen = ax( + f() + .input('topic', z.string().min(1)) + .output('summary', z.string().min(50)) + .build() +); + +// bestOfN: choose the best complete candidate. +const selected = bestOfN(gen, { + n: 4, + rewardFn: ({ prediction }) => prediction.summary.length, +}); + +// Whole-output assertion: retries with correction feedback. +gen.addAssert( + (output) => output.summary.includes(topic) || 'Summary must mention the topic.' +); + +// Streaming assertion: fail fast on unsafe partial output. +gen.addStreamingAssert( + 'summary', + (text) => !text.includes('forbidden'), + 'Output contains forbidden text' +); +``` + +Rules: + +- Schema validation retries with parser/constraint feedback. +- `addAssert(...)` checks the complete parsed output after validation/processors and retries with correction feedback on failure. +- `bestOfN(...)` scores complete candidates and returns the highest reward or first threshold hit. +- `refine(...)` runs rounds and can feed reward-derived advice into instruction components between rounds. +- `addStreamingAssert(...)` targets a string/code output field and receives partial text so far. +- Streaming assertions abort the current stream attempt by throwing `AxStreamingAssertionError`, then feed correction feedback into AxGen retries. + +## Field Processors + +```typescript +// Post-processing after generation +gen.addFieldProcessor('summary', (value, context) => value.toUpperCase()); + +// Streaming field processor (called on each chunk) +gen.addStreamingFieldProcessor('content', (partialValue, context) => { + console.log(`Received ${partialValue.length} chars`); + return partialValue; +}); +``` + +Rules: + +- `addFieldProcessor` runs once after the field is fully generated. +- `addStreamingFieldProcessor` runs on each streaming chunk for the target field. +- Both must return the (possibly transformed) value. + +## Function Calling + +```typescript +const result = await gen.forward(llm, { question: '...' }, { + functions: tools, + functionCallMode: 'auto', + stopFunction: 'finalAnswer', +}); +``` + +Rules: + +- `functionCallMode` can be `'auto'`, `'none'`, or a specific function name to force. +- `stopFunction` accepts a string or string[] to halt multi-step on specific function calls. +- Multi-step continues until all outputs filled, stop function called, or `maxSteps` reached. + +## Caching + +### Response Caching + +```typescript +const gen = ax('question:string -> answer:string', { + cachingFunction: async (key, value?) => { + if (value !== undefined) { + await cache.set(key, value); + return; + } + return await cache.get(key); + }, +}); +``` + +### Context Caching + +```typescript +const result = await gen.forward(llm, { question: '...' }, { + contextCache: { cacheBreakpoint: 'after-examples' }, +}); +``` + +Rules: + +- `cachingFunction` acts as a get/set: called with `(key)` to read, `(key, value)` to write. +- `contextCache` enables AI provider-level prompt caching for long context. + +## Sampling And Result Picker + +```typescript +const result = await gen.forward(llm, { question: '...' }, { + sampleCount: 3, + resultPicker: async (samples) => { + // Evaluate each sample and return the index of the best one + return bestIndex; + }, +}); +``` + +Rules: + +- `sampleCount` generates multiple completions in parallel. +- `resultPicker` receives all samples and must return the index of the chosen result. + +## Extended Thinking + +```typescript +const result = await gen.forward(llm, { question: '...' }, { + thinkingTokenBudget: 'medium', + showThoughts: true, +}); +console.log(result.thought); +``` + +Rules: + +- `thinkingTokenBudget` can be `'low'`, `'medium'`, `'high'`, or a number. +- Set `showThoughts: true` to include the model's reasoning in `result.thought`. + +## Structured Outputs + +```typescript +const sig = f() + .input('text', f.string()) + .output('summary', f.string()) + .output('metadata', f.json().optional()) + .useStructured() + .build(); +``` + +Rules: + +- `.useStructured()` asks providers with native support, including OpenAI, Anthropic, and Gemini, for schema-constrained JSON. +- Native structured-output schemas list every object property in `required`, set `additionalProperties: false` on objects, and express optional fields as nullable types. +- Flexible `json` fields and unshaped `object` fields are sent as JSON-encoded strings for native structured outputs, then parsed back into normal JavaScript values. + +## Step Hooks + +```typescript +const result = await gen.forward(llm, values, { + stepHooks: { + beforeStep: (ctx) => { + if (ctx.functionsExecuted.has('complexanalysis')) { + ctx.setModel('smart'); + ctx.setThinkingBudget('high'); + } + }, + afterStep: (ctx) => { + console.log(`Usage: ${ctx.usage.totalTokens} tokens`); + }, + }, +}); +``` + +### AxStepContext Read-Only Properties + +- `stepIndex` - current step number +- `maxSteps` - configured maximum steps +- `isFirstStep` - whether this is the first step +- `functionsExecuted` - `Set` of function names called so far +- `lastFunctionCalls` - array of the most recent function call results +- `usage` - token usage statistics +- `state` - current step state + +### AxStepContext Mutators + +- `setModel(model)` - change the model for the next step +- `setThinkingBudget(budget)` - adjust thinking budget +- `setTemperature(temp)` - adjust temperature +- `setMaxTokens(max)` - adjust max output tokens +- `setOptions(opts)` - set arbitrary forward options +- `addFunctions(fns)` - add functions for the next step +- `removeFunctions(names)` - remove functions by name +- `stop()` - stop multi-step execution + +Rules: + +- All mutations are pending and applied at the next step boundary. +- `beforeStep` runs before each LLM call; `afterStep` runs after. +- Use `afterFunctionExecution` to react to specific function results. + +## Self-Tuning + +```typescript +// Simple: enable all self-tuning +const result = await gen.forward(llm, values, { selfTuning: true }); + +// Granular: pick what to tune +const result = await gen.forward(llm, values, { + selfTuning: { + model: true, + thinkingBudget: true, + functions: [searchWeb, calculate], + }, +}); +``` + +Rules: + +- `selfTuning: true` enables automatic model and parameter selection. +- Granular config allows tuning specific aspects independently. +- `selfTuning.functions` provides a pool of functions the tuner may add or remove per step. + +## Error Handling + +```typescript +import { AxGenerateError } from '@ax-llm/ax'; + +try { + const result = await gen.forward(llm, { input: '...' }); +} catch (error) { + if (error instanceof AxGenerateError) { + console.log(error.details.model, error.details.signature); + } +} +``` + +Rules: + +- `AxGenerateError` includes `details` with `model` and `signature` for debugging. +- `AxAIServiceAbortedError` is thrown on cancellation via `stop()` or `abortSignal`. + +## Chat Log and Usage + +### getChatLog() + +After any `.forward()` or `streamingForward()` call, `gen.getChatLog()` returns the full normalized chat history — every `ai.chat()` round-trip, including the system prompt, all messages, and the model response. The log is reset at the start of each `.forward()` call. Multi-step generators (with function calls) produce one entry per step. + +```typescript +await gen.forward(llm, { question: 'What is 2+2?' }); + +for (const entry of gen.getChatLog()) { + console.log('model:', entry.model); + for (const msg of entry.messages) { + console.log(`[${msg.role}]`, msg.content); + } + console.log('tokens:', entry.modelUsage?.tokens); +} +``` + +Message roles: `system`, `user`, `assistant`, `tool`. Assistant content uses inline XML: +- `...` — reasoning/thinking tokens +- `\n{...}\n` — tool invocations + +The system message includes a `` JSON block when functions are present. + +```typescript +type AxChatLogMessage = + | { role: 'system'; content: string } + | { role: 'user'; content: string } + | { role: 'assistant'; content: string } + | { role: 'tool'; name: string; content: string }; + +type AxChatLogEntry = { + name?: string; + model: string; + messages: AxChatLogMessage[]; + modelUsage?: AxProgramUsage; +}; + +gen.getChatLog(): readonly AxChatLogEntry[] +``` + +### getUsage() + +Returns token usage aggregated by `(ai, model)` across all steps. When a provider reports prompt-cache usage, `promptTokens` is the uncached input portion and `cacheReadTokens` / `cacheCreationTokens` carry the cache counters. Reset with `resetUsage()`. + +```typescript +const usage = gen.getUsage(); // AxProgramUsage[] +console.log(usage[0]?.tokens?.promptTokens); +gen.resetUsage(); +``` + +`AxAgent` and `AxFlow` also return flat `AxChatLogEntry[]` logs; composite programs set `entry.name` so callers can filter by node/stage. + +## Examples + +Fetch these for full working code: + +- [Streaming](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/streaming.ts) — field-by-field streaming +- [Best Of N](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/best-of-n.ts) — reward-scored sample selection +- [Refine](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/refine.ts) — retry rounds with generated feedback +- [Streaming Assert](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/streaming-asserts.ts) — fail-fast partial-output correction +- [Structured Output](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/structured_output.ts) — fluent API with validation +- [Debug Logging](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/debug-logging.ts) — debug mode and step hooks +- [Stop Function](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/stop-function.ts) — stop functions +- [Fibonacci](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/fibonacci.ts) — streaming with thinking +- [Extraction](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/extract.ts) — information extraction +- [Multi-Sampling](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/sample-count.ts) — sample count usage + +## Do Not Generate + +- Do not use `new AxGen(...)` for new code unless explicitly required. +- Do not pass raw API keys or config objects where an `ai(...)` instance is expected. +- Do not use `forward()` for streaming; use `streamingForward()`. +- Do not use streaming assertions as reward/refine mechanisms; they enforce hard partial-output invariants and retry with correction. +- Do not mutate step hook context expecting immediate effect; mutations are pending until the next step. +- Do not assume multi-step stops after one LLM call; it continues until outputs are filled, a stop function fires, or `maxSteps` is reached. diff --git a/.claude/skills/ax-gepa/SKILL.md b/.claude/skills/ax-gepa/SKILL.md new file mode 100644 index 0000000..44530ca --- /dev/null +++ b/.claude/skills/ax-gepa/SKILL.md @@ -0,0 +1,264 @@ +--- +name: ax-gepa +description: This skill helps an LLM generate correct AxGEPA optimization code using @ax-llm/ax. Use when the user asks about AxGEPA, GEPA, Pareto optimization, multi-objective prompt tuning, reflective prompt evolution, validationExamples, maxMetricCalls, or optimizing a generator, flow, or agent tree. +version: "23.0.0" +--- + +# GEPA Optimization Codegen Rules (@ax-llm/ax) + +Use this skill to generate GEPA optimization code. Prefer the top-level `optimize(...)` helper for normal code, and use direct `AxGEPA` / `AxBootstrapFewShot` only when the user needs low-level optimizer control. + +## Use These Defaults + +- Use `optimize(program, train, metric, { studentAI, teacherAI, ... })` for normal generator and flow tuning. +- Prefer `ai()`, `ax()`, and `flow()` for new code. +- Use a strong `teacherAI` and a cheaper `studentAI`. +- Pass `validationExamples` when you have a holdout set. +- Set `maxMetricCalls` to bound optimizer cost; `optimize(...)` defaults it to `100`. +- Use scalar metrics for one objective and object metrics for Pareto optimization. +- Apply results with `program.applyOptimization(result.optimizedProgram!)`. +- For tree-wide runs, expect `optimizedProgram.componentMap`. +- Persist artifacts with `axSerializeOptimizedProgram(...)` and restore them with `axDeserializeOptimizedProgram(...)` so the same flow works in browsers and Node. +- `optimize(...)` runs `AxBootstrapFewShot -> AxGEPA` for small starter sets by default, preserving the demos in `result.optimizedProgram.demos`. + +## Critical Rules + +- `optimize(...)` and `AxGEPA.compile()` work for a single generator and for tree-aware roots such as flows or agents with registered optimizable descendants. +- There is no separate flow-only GEPA optimizer. Use `AxGEPA` for flows too. +- The metric may return either `number` or `Record`. +- Keep metrics deterministic and cheap by default. +- Avoid extra LLM calls inside the metric unless the user explicitly wants judge-based evaluation. +- If the user needs LLM-as-judge scoring for a non-agent GEPA run, prefer a plain typed `AxGen` evaluator instead of writing a custom judge abstraction. +- `maxMetricCalls` must be large enough to cover the initial validation pass over `validationExamples`. +- GEPA optimizes generic string components exposed by `getOptimizableComponents()`. If a tree exposes no components, optimization will fail. +- Use held-out validation examples for selection. Do not reuse the training set as `validationExamples`. +- `result.optimizedProgram` is the easy-to-apply best candidate. `result.paretoFront` is the full trade-off set for multi-objective runs. +- Direct `AxGEPA` still has its own `bootstrap` option, but top-level `optimize(...)` composes the existing `AxBootstrapFewShot` optimizer before GEPA instead. + +## Metric Selection + +Choose the evaluation path deliberately: + +- Prefer a deterministic metric when correctness can be read directly from `prediction` and `example`. +- Prefer a deterministic metric when cost, latency, recursion depth, or tool count matters. +- Use a plain typed `AxGen` evaluator only when the task is genuinely qualitative and hard to score exactly. +- For `agent.optimize(...)`, prefer the built-in judge path instead of manually wrapping a judge metric. Normal agent users usually do not need to set `target` or `metric` at all. + +Rule of thumb: + +- `optimize(...)` on `AxGen` or flow: use a metric first, optionally a plain typed `AxGen` evaluator if needed. +- `agent.optimize(...)`: use custom `metric` for crisp scoring, otherwise let the built-in judge handle scoring. Add `judgeAI` plus `judgeOptions` only when you want a stronger or separate judge model. + +## Canonical Scalar Pattern + +```typescript +import { ai, ax, optimize, AxAIOpenAIModel } from '@ax-llm/ax'; + +const student = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, + config: { model: AxAIOpenAIModel.GPT54Mini }, +}); + +const teacher = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, + config: { model: AxAIOpenAIModel.GPT54 }, +}); + +const classifier = ax( + 'emailText:string -> priority:class "high, normal, low", rationale:string' +); + +const train = [ + { emailText: 'URGENT: Server down!', priority: 'high' }, + { emailText: 'Weekly newsletter', priority: 'low' }, +]; + +const validation = [ + { emailText: 'Invoice overdue', priority: 'high' }, + { emailText: 'Lunch plans?', priority: 'low' }, +]; + +const metric = ({ prediction, example }: { prediction: any; example: any }) => + prediction?.priority === example?.priority ? 1 : 0; + +const result = await optimize(classifier, train, metric, { + studentAI: student, + teacherAI: teacher, + numTrials: 12, + minibatch: true, + minibatchSize: 4, + earlyStoppingTrials: 4, + sampleCount: 1, + validationExamples: validation, + maxMetricCalls: 120, +}); + +classifier.applyOptimization(result.optimizedProgram!); +console.log(result.bestScore); +``` + +## Canonical Pareto Pattern + +```typescript +import { ai, flow, optimize, AxAIOpenAIModel } from '@ax-llm/ax'; + +const student = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, + config: { model: AxAIOpenAIModel.GPT54Mini }, +}); + +const teacher = ai({ + name: 'openai', + apiKey: process.env.OPENAI_APIKEY!, + config: { model: AxAIOpenAIModel.GPT54 }, +}); + +const wf = flow<{ emailText: string }>() + .n('classifier', 'emailText:string -> priority:class "high, normal, low"') + .n( + 'rationale', + 'emailText:string, priority:string -> rationale:string "One concise sentence"' + ) + .e('classifier', (state) => ({ emailText: state.emailText })) + .e('rationale', (state) => ({ + emailText: state.emailText, + priority: state.classifierResult.priority, + })) + .r((state) => ({ + priority: state.classifierResult.priority, + rationale: state.rationaleResult.rationale, + })); + +const train = [ + { emailText: 'URGENT: Server down!', priority: 'high' }, + { emailText: 'Weekly newsletter', priority: 'low' }, +]; + +const validation = [ + { emailText: 'Invoice overdue', priority: 'high' }, + { emailText: 'Lunch plans?', priority: 'low' }, +]; + +const metric = ({ prediction, example }: { prediction: any; example: any }) => { + const accuracy = prediction?.priority === example?.priority ? 1 : 0; + const rationale = typeof prediction?.rationale === 'string' + ? prediction.rationale + : ''; + const brevity = rationale.length <= 40 ? 1 : rationale.length <= 80 ? 0.5 : 0.1; + return { accuracy, brevity }; +}; + +const result = await optimize(wf, train, metric, { + studentAI: student, + teacherAI: teacher, + numTrials: 16, + minibatch: true, + minibatchSize: 6, + earlyStoppingTrials: 5, + sampleCount: 1, + validationExamples: validation, + maxMetricCalls: 240, +}); + +for (const point of result.paretoFront) { + console.log(point.scores, point.configuration); +} + +wf.applyOptimization(result.optimizedProgram!); +console.log(result.optimizedProgram?.componentMap); +``` + +## Metric Patterns + +```typescript +// Scalar objective +const scalarMetric = ({ prediction, example }) => + prediction.answer === example.answer ? 1 : 0; + +// Multi-objective +const multiMetric = ({ prediction, example }) => ({ + accuracy: prediction.answer === example.answer ? 1 : 0, + brevity: + typeof prediction?.reasoning === 'string' && + prediction.reasoning.length < 120 + ? 1 + : 0.2, +}); +``` + +- Return plain numbers or plain object literals. +- Keep objective names stable across calls. +- Prefer normalized scores such as `0..1` so trade-offs are easy to reason about. + +## Result Handling + +```typescript +const { optimizedProgram, paretoFront } = result; + +program.applyOptimization(optimizedProgram!); + +// Save for later +const saved = JSON.stringify(optimizedProgram); + +// Load later and re-apply +const loaded = JSON.parse(saved); +program.applyOptimization(loaded); +``` + +- Single-target runs usually populate both `optimizedProgram.instruction` and `optimizedProgram.componentMap`. +- Tree-wide runs rely on `componentMap`, keyed by full component key. +- Pareto points expose candidate configs under `point.configuration.componentMap`. + +## Useful Options + +```typescript +const optimizer = new AxGEPA({ + studentAI, + teacherAI, + numTrials: 20, + minibatch: true, + minibatchSize: 5, + minibatchFullEvalSteps: 5, + earlyStoppingTrials: 5, + minImprovementThreshold: 0, + sampleCount: 1, + seed: 42, + verbose: true, +}); +``` + +- `numTrials`: number of reflection/evolution rounds. +- `minibatch`: reduce per-round evaluation cost. +- `minibatchSize`: examples per minibatch. +- `earlyStoppingTrials`: stop after repeated non-improvement. +- `minImprovementThreshold`: reject tiny gains below this threshold. +- `seed`: stabilize sampling during demos and tests. + +## Budgeting and Validation + +- Always create distinct `train` and `validationExamples` arrays. +- Size `maxMetricCalls` for at least one full validation pass plus several rounds. +- If the user wants a strict budget, say so explicitly and set `maxMetricCalls`. +- For expensive trees, start with `auto: 'light'` or fewer `numTrials`, then scale up. +- GEPA selects among exposed components using measured accept/reject history, not LLM-generated numeric scores. The LLM proposes component text; metrics decide whether to keep it. +- Function/tool trace reflection is keyed by stable component IDs where available, so function renames do not break saved candidate maps. + +## Troubleshooting + +- Error about `maxMetricCalls` being too small: increase it until the initial validation pass fits. +- Empty or poor Pareto front: verify the metric returns numbers for every example. +- No tree optimization effect: ensure child programs are registered under the root and expose optimizable components. +- Saved optimization applies only partly: use `program.applyOptimization(...)`, not just `setInstruction(...)`, so `componentMap` reaches the full tree. +- Agent target seems too broad: when using `agent.optimize(...)`, set `target: 'actor'`, `'responder'`, `'all'`, or explicit program IDs. The wrapper filters GEPA components to the selected target. + +## Good Example Targets + +- `/Users/vr/src/ax/src/examples/optimize.ts` +- `/Users/vr/src/ax/src/examples/gepa.ts` +- `/Users/vr/src/ax/src/examples/gepa-flow.ts` +- `/Users/vr/src/ax/src/examples/gepa-train-inference.ts` +- `/Users/vr/src/ax/src/examples/gepa-quality-vs-speed-optimization.ts` +- `/Users/vr/src/ax/src/examples/axagent-gepa-optimization.ts` diff --git a/.claude/skills/ax-llm/SKILL.md b/.claude/skills/ax-llm/SKILL.md new file mode 100644 index 0000000..998d698 --- /dev/null +++ b/.claude/skills/ax-llm/SKILL.md @@ -0,0 +1,339 @@ +--- +name: ax-llm +description: This skill helps with using the @ax-llm/ax TypeScript library for building LLM applications. Use when the user asks about ax(), ai(), f(), s(), agent(), flow(), AxGen, AxAgent, AxFlow, signatures, streaming, or mentions @ax-llm/ax. +version: "23.0.0" +--- + +# Ax Library (@ax-llm/ax) Quick Reference + +Ax is a TypeScript library for building LLM-powered applications with type-safe signatures, streaming support, and multi-provider compatibility. + +> **Detailed skills available:** ax-ai (providers), ax-signature (signatures/types), ax-gen (generators), ax-agent (core agents/tools), ax-agent-rlm (agent runtime/RLM/delegation), ax-agent-observability (callbacks/logs/usage), ax-agent-memory-skills (recall and dynamic skill loading), ax-agent-optimize (agent tuning/eval), ax-flow (workflows), ax-gepa (top-level `optimize(...)`, BootstrapFewShot -> GEPA, Pareto optimization). + +## Imports & Factories + +```typescript +// Prefer factory functions: ax(), ai(), agent(), flow(); avoid class constructors. +import { ax, ai, f, s, fn, agent, flow, AxMemory, AxMCPClient } from '@ax-llm/ax'; +import { z } from 'zod'; // optional — any Standard Schema v1 library works + +// AI provider +const llm = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY }); + +// Generator (from string signature) +const gen = ax('question:string -> answer:string'); + +// Generator (from fluent signature) +const gen = ax( + f() + .input('question', f.string('User question')) + .output('answer', f.string('AI response')) + .build() +); + +// Generator (from zod — Standard Schema v1, also works with valibot/arktype) +const zodGen = ax( + f() + .input(z.object({ question: z.string().describe('User question') })) + .output(z.object({ answer: z.string().describe('AI response') })) + .build() +); + +// Reusable signature +const sig = s('question:string, context:string[] -> answer:string'); + +// Agent +const myAgent = agent('userInput:string -> response:string', { + name: 'helper', + description: 'A helpful assistant', +}); + +// Flow +const wf = flow<{ input: string }, { output: string }>() + .node('step1', 'input:string -> output:string') + .execute('step1', (state) => ({ input: state.input })) + .returns((state) => ({ output: state.step1Result.output })); + +// Function tool — native fluent +const tool = fn('search') + .description('Search the web') + .arg('query', f.string('Search query')) + .returns(f.string('Search results')) + .handler(({ query }) => searchWeb(query)) + .build(); + +// Function tool — zod schema (Standard Schema v1: also works with valibot, arktype) +const zodTool = fn('calculateTax') + .description('Calculate tax for an amount') + .arg(z.object({ + amount: z.number().positive().describe('Pre-tax amount in USD'), + region: z.enum(['US', 'EU', 'UK']).describe('Tax region'), + })) + .returns(z.object({ tax: z.number(), total: z.number() })) + .handler(async ({ amount }) => ({ tax: amount * 0.1, total: amount * 1.1 })) + .build(); +``` + +## Running + +```typescript +// Forward (blocking) +const result = await gen.forward(llm, { question: 'What is 2+2?' }); + +// Streaming +for await (const chunk of gen.streamingForward(llm, { question: 'Tell a story' })) { + if (chunk.delta.answer) process.stdout.write(chunk.delta.answer); +} +``` + +## Forward Options Quick Reference + +| Goal | Option | Example | +|------|--------|---------| +| Model override | `model` | `{ model: 'gpt-5.4-mini' }` | +| Temperature | `modelConfig.temperature` | `{ modelConfig: { temperature: 0.8 } }` | +| Max tokens | `modelConfig.maxTokens` | `{ modelConfig: { maxTokens: 500 } }` | +| Retry on failure | `maxRetries` | `{ maxRetries: 3 }` | +| Max agent steps | `maxSteps` | `{ maxSteps: 10 }` | +| Fail fast | `fastFail` | `{ fastFail: true }` | +| Thinking budget | `thinkingTokenBudget` | `{ thinkingTokenBudget: 'medium' }` | +| Show thoughts | `showThoughts` | `{ showThoughts: true }` | +| Context caching | `contextCache` | `{ contextCache: { cacheBreakpoint: 'after-examples' } }` | +| Multi-sampling | `sampleCount` | `{ sampleCount: 5 }` | +| Debug logging | `debug` | `{ debug: true }` | +| Abort signal | `abortSignal` | `{ abortSignal: controller.signal }` | +| Memory | `mem` | `{ mem: new AxMemory() }` | +| Stop function | `stopFunction` | `{ stopFunction: 'finalAnswer' }` | +| Function mode | `functionCallMode` | `{ functionCallMode: 'auto' }` | + +Global runtime defaults can be set with `axGlobals` and are read live by future AI, AxGen, and AxFlow calls: + +```typescript +import { axGlobals, axCreateDefaultColorLogger } from '@ax-llm/ax'; +import { trace } from '@opentelemetry/api'; + +axGlobals.tracer = trace.getTracer('my-app'); +axGlobals.debug = true; +axGlobals.logger = axCreateDefaultColorLogger(); +``` + +Precedence is: per-call options, then explicit instance/program options, then current `axGlobals`, then built-in defaults. `customLabels` merge in that order, and `abortSignal` values are combined so either global or local cancellation works. + +## Memory and Context + +```typescript +import { AxMemory } from '@ax-llm/ax'; + +const memory = new AxMemory(); + +// Multi-turn conversation +await gen.forward(llm, { userMessage: 'My name is Alice' }, { mem: memory }); +const r = await gen.forward(llm, { userMessage: 'What is my name?' }, { mem: memory }); +``` + +## Few-Shot Examples + +```typescript +const classifier = ax('reviewText:string -> sentiment:class "positive, negative, neutral"'); + +classifier.setExamples([ + { reviewText: 'I love this!', sentiment: 'positive' }, + { reviewText: 'Terrible.', sentiment: 'negative' }, + { reviewText: 'It works.', sentiment: 'neutral' }, +]); +``` + +## Common Patterns + +### Classification + +```typescript +const classifier = ax( + f() + .input('text', f.string()) + .output('category', f.class(['spam', 'ham', 'uncertain'])) + .output('confidence', f.number().min(0).max(1)) + .build() +); +``` + +### Extraction + +```typescript +const extractor = ax( + f() + .input('text', f.string()) + .output('entities', f.object({ + people: f.string().array(), + organizations: f.string().array(), + locations: f.string().array() + })) + .build() +); +``` + +### Multi-modal (Images) + +```typescript +const analyzer = ax( + f() + .input('image', f.image('Image to analyze')) + .input('question', f.string('Question').optional()) + .output('description', f.string()) + .output('objects', f.string().array()) + .build() +); + +const result = await analyzer.forward(llm, { + image: { mimeType: 'image/jpeg', data: base64Data }, + question: 'What objects are in this image?' +}); +``` + +### Chaining Generators + +```typescript +const researcher = ax('topic:string -> research:string, keyFacts:string[]'); +const writer = ax('research:string, keyFacts:string[] -> article:string'); + +const research = await researcher.forward(llm, { topic: 'AGI' }); +const draft = await writer.forward(llm, { research: research.research, keyFacts: research.keyFacts }); +``` + +## Error Handling + +```typescript +import { AxGenerateError, AxAIServiceError, AxAIServiceAbortedError } from '@ax-llm/ax'; + +try { + const result = await gen.forward(llm, { input: 'test' }); +} catch (error) { + if (error instanceof AxGenerateError) { + console.error('Generation failed:', error.details.model, error.details.signature); + } else if (error instanceof AxAIServiceAbortedError) { + console.log('Request was aborted'); + } else if (error instanceof AxAIServiceError) { + console.error('AI service error:', error.message); + } +} +``` + +## Debugging + +```typescript +import { axCreateDefaultColorLogger, axGlobals } from '@ax-llm/ax'; + +const result = await gen.forward(llm, { input: 'test' }, { + debug: true, + logger: axCreateDefaultColorLogger(), + // OpenTelemetry + tracer: openTelemetryTracer, + meter: openTelemetryMeter, +}); + +// Or set live app-wide defaults for future calls: +axGlobals.tracer = openTelemetryTracer; +axGlobals.meter = openTelemetryMeter; +``` + +## MCP Integration + +```typescript +import { AxMCPClient, agent } from '@ax-llm/ax'; +import { AxMCPStdioTransport } from '@ax-llm/ax-tools'; + +// Stdio transport (local MCP server) +const transport = new AxMCPStdioTransport({ + command: 'npx', + args: ['-y', '@modelcontextprotocol/server-memory'], +}); + +const mcpClient = new AxMCPClient(transport, { debug: false }); +await mcpClient.init(); + +// Use with agent under a namespace +const myAgent = agent('userMessage:string -> response:string', { + functions: [ + { + namespace: 'memory', + title: 'Memory MCP', + description: 'Memory server tools', + selectionCriteria: 'Use for persistent memory lookup and updates.', + functions: [mcpClient], + }, + ], + functionDiscovery: true, + contextFields: [], +}); +``` + +### HTTP Transport (Remote MCP) + +```typescript +import { AxMCPStreambleHTTPTransport } from '@ax-llm/ax/mcp/transports/httpStreamTransport.js'; + +const transport = new AxMCPStreambleHTTPTransport('https://remote.mcp.pipedream.net', { + headers: { 'x-pd-project-id': projectId }, + authorization: `Bearer ${accessToken}`, +}); +``` + +### MCP Capabilities + +| Capability | Prefix | Description | +|---|---|---| +| Tools | *(none)* | Function calls | +| Prompts | `prompt_` | Prompt templates | +| Resources | `resource_` | File/data access | + +```typescript +const caps = mcpClient.getCapabilities(); +const functions = mcpClient.toFunction(); +``` + +### Function Overrides + +```typescript +const mcpClient = new AxMCPClient(transport, { + functionOverrides: [ + { name: 'search_documents', updates: { name: 'findDocs', description: 'Search docs' } } + ] +}); +``` + +## Type Reference + +```typescript +class AxGen { + forward(ai: AxAIService, values: IN, options?: AxProgramForwardOptions): Promise; + streamingForward(ai: AxAIService, values: IN, options?: AxProgramStreamingForwardOptions): AsyncGenerator<{ delta: Partial }>; + setExamples(examples: Array>): void; + addAssert(fn: (output: OUT) => boolean | string | undefined | Promise, message?: string): void; + addStreamingAssert(field: keyof OUT, fn: (chunk: string, done?: boolean) => boolean | string | undefined | Promise, message?: string): void; + addFieldProcessor(field: keyof OUT, fn: (value: any) => any): void; + addStreamingFieldProcessor(field: keyof OUT, fn: (chunk: string, ctx: any) => void): void; + stop(): void; +} + +class AxAgent { + forward(ai: AxAIService, values: IN, options?: AxAgentOptions): Promise; + streamingForward(ai: AxAIService, values: IN, options?: AxAgentOptions): AsyncGenerator<{ delta: Partial }>; + getFunction(): AxFunction; +} + +class AxFlow { + node(name: string, signature: string | AxSignature): AxFlow; + execute(name: string, mapper: (state) => any): AxFlow; + returns(mapper: (state) => OUT): AxFlow; + forward(ai: AxAIService, values: IN): Promise; +} +``` + +## Examples + +Fetch these for full working code: + +- [Standard Schema (zod)](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/standard-schema.ts) — zod with f() and fn() +- [Chat](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/chat.ts) — multi-turn conversation +- [Marketing](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/marketing.ts) — product use case +- [MCP Integration](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/mcp-client-memory.ts) — MCP integration diff --git a/.claude/skills/ax-playbook/SKILL.md b/.claude/skills/ax-playbook/SKILL.md new file mode 100644 index 0000000..c1b86f0 --- /dev/null +++ b/.claude/skills/ax-playbook/SKILL.md @@ -0,0 +1,95 @@ +--- +name: ax-playbook +description: This skill helps an LLM generate correct playbook code using @ax-llm/ax. Use when the user asks about playbook(), AxPlaybook, context playbooks, evolving context, ACE / Agentic Context Engineering, agent.playbook(), or growing/applying task knowledge offline and online with evolve() and update(). +version: "23.0.0" +--- + +# Playbook Codegen Rules (@ax-llm/ax) + +Use this skill to generate context-playbook code. A playbook grows an evolving body of task knowledge and renders it into a program's context. The evolution engine (ACE — Agentic Context Engineering) is hidden behind `playbook(...)`, exactly as `optimize(...)` hides its optimizer. Prefer the `playbook(...)` concept; only reach for `AxACE` directly when the user explicitly wants the low-level engine. + +## Use These Defaults + +- Create with `playbook(program, { studentAI, teacherAI? })`; it returns an `AxPlaybook` handle. +- Grow offline with `await pb.evolve(examples, metric)` — returns `{ bestScore, playbook }`. +- Grow online with `await pb.update({ example, prediction, feedback })` — no metric needed. +- Apply with `pb.applyTo(program)` (defaults to the bound program). +- Persist with `pb.toJSON()` and restore with `playbook(program, opts).load(snapshot)`. +- Inspect with `pb.render()` (markdown) and `pb.getState()` (`{ playbook, artifact }`). +- For agents use `agent.playbook({ target: 'actor' | 'responder' })`; default target is `'actor'`. +- Use a cheaper `studentAI` to run the program and an optional stronger `teacherAI` to reflect/curate. +- Prefer `ai()`, `ax()`, and `agent()` for new code. + +## Critical Rules + +- `playbook(...)` binds to an `AxGen` program; `evolve`/`update` need that program's signature. +- `evolve()` returns only `{ bestScore, playbook }`. There is no Pareto front and no `optimizedProgram` — that is `optimize(...)`'s shape, not a playbook's. +- `update({ example, prediction, feedback })` requires the full `{ example, prediction }`; `example` must match the program's input fields (plus any expected output). Do not pass bare input fields at the top level. +- `update()` works without a prior `evolve()`/`load()` — the handle hydrates lazily on first use. +- `applyTo()` injects a `## Context Playbook` block into the program description; calling it repeatedly recomposes from the original base (no stacking). +- Keep the offline `metric` deterministic and cheap, like a GEPA metric. +- A playbook is plain JSON. Persist `pb.toJSON()` and `load(...)` it into a fresh program for production. +- This is a TypeScript feature; do not suggest it for the generated (Python/Go/Rust/Java/C++) packages yet. + +## Offline Pattern (evolve) + +```typescript +import { type AxMetricFn, ai, ax, playbook } from '@ax-llm/ax'; + +const program = ax('review:string -> sentiment:class "positive, negative"'); +const studentAI = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY! }); +const metric: AxMetricFn = ({ prediction, example }) => + (prediction as any).sentiment === (example as any).sentiment ? 1 : 0; + +const pb = playbook(program, { studentAI, maxEpochs: 2 }); +const { bestScore } = await pb.evolve(train, metric); +pb.applyTo(program); +``` + +## Online Pattern (update) + +```typescript +// After a real run, feed the outcome back so the playbook keeps learning. +await pb.update({ + example: { review: 'Five stars, would buy again.' }, + prediction: { sentiment: 'negative' }, + feedback: 'WRONG: enthusiastic praise is positive.', +}); +pb.applyTo(program); +``` + +## Persist And Restore + +```typescript +const snapshot = pb.toJSON(); // { playbook, artifact } — plain JSON +// later, in another process / a production program instance: +playbook(prodProgram, { studentAI }).load(snapshot).applyTo(prodProgram); +``` + +## Agents + +```typescript +const a = agent('ticket:string -> reply:string', { ai }); +const apb = a.playbook({ target: 'actor' }); // 'actor' (default) or 'responder' +await apb.update({ example, prediction, feedback }); // injected into the live stage prompt +``` + +Offline `evolve(...)` on an agent stage scores that stage in isolation; for full-pipeline tuning of agent instructions and demos use `agent.optimize(...)` (GEPA). + +## Playbook vs optimize() + +- `playbook(...)` — accumulate reusable, evolving task knowledge; the only path that also learns online via `update(...)`. +- `optimize(...)` / `agent.optimize(...)` — tune instruction text and few-shot demos offline to a best/Pareto result. +- They are complementary; a project can use both. + +## Troubleshooting + +- "Cannot convert undefined or null to object" from `update()` → you passed input fields at the top level; wrap them in `example: { ... }`. +- Empty playbook after `evolve()` → the model already scored well, so nothing was curated; use harder/ambiguous examples or a weaker `studentAI` to surface lessons. +- Playbook not affecting an agent's behavior → ensure `apply` is not `false` and you used `agent.playbook(...)` (not a bare `playbook()` on an internal program). + +## See Also + +- `ax-gepa` - `optimize(...)` and `AxGEPA` for instruction/demo tuning. +- `ax-agent-context` - choosing between contextMap, contextPolicy, `agent.playbook(...)`, and recall. +- `ax-agent-optimize` - `agent.optimize(...)` GEPA tuning for agents. diff --git a/.claude/skills/ax-refine/SKILL.md b/.claude/skills/ax-refine/SKILL.md new file mode 100644 index 0000000..b7ef0ac --- /dev/null +++ b/.claude/skills/ax-refine/SKILL.md @@ -0,0 +1,81 @@ +--- +name: ax-refine +description: Use this skill when writing or reviewing Ax bestOfN/refine code, reward functions, thresholds, native sample selection, serial attempts, generated advice, and attempt diagnostics. +version: "23.0.0" +--- + +# Ax Refine And BestOfN + +Use `bestOfN(...)` when you can score complete outputs independently. Use `refine(...)` when failed rounds should produce feedback that changes the next attempt. + +## Validation And Assertions + +Keep reward scoring, whole-output assertions, and streaming assertions separate: + +- Use schema validation for shape, types, and field-level constraints. +- Use `addAssert(...)` for whole-output hard invariants. Failed assertions feed correction text into the normal retry loop. +- Use `addStreamingAssert(...)` for partial streaming hard invariants. It aborts the current stream attempt as soon as the partial field fails, then feeds correction text into the normal retry loop. +- Use `bestOfN(...)` for complete-candidate selection. +- Use `refine(...)` for reward-scored retry rounds with generated feedback. + +## APIs + +```typescript +import { bestOfN, refine } from '@ax-llm/ax'; + +const selected = bestOfN(program, { + n: 4, + threshold: 0.8, + rewardFn: ({ input, prediction, traces, chatLog }) => score(prediction), +}); + +const improved = refine(program, { + rounds: 3, + samplesPerRound: 2, + threshold: 0.85, + rewardDescription: 'Prefer complete, grounded, concise answers.', + rewardFn: ({ prediction }) => score(prediction), +}); +``` + +Rules: + +- `forward(...)` returns the selected prediction. +- `streamingForward(...)` is unsupported; score complete outputs instead. +- `getUsage()` aggregates usage across attempts. +- `getTraces()` and `getChatLog()` return the selected attempt's diagnostics. +- `getAttempts()` returns all attempt metadata, including reward, errors, and advice application. + +## Reward Functions + +Reward functions return a number. Higher is better. A `threshold` marks a good-enough candidate and can stop serial attempts early. + +```typescript +const rewardFn = ({ prediction }) => { + const exact = prediction.answer === 'Paris' ? 1 : 0; + const concise = prediction.answer.length < 80 ? 0.2 : 0; + return exact + concise; +}; +``` + +Use serial strategy when the reward needs traces, chat logs, tools, or full flow behavior. + +## Strategies + +- `strategy: "auto"` uses native samples for `AxGen` and serial attempts for composite programs. +- `strategy: "native-samples"` uses `sampleCount` and a reward-backed `resultPicker`; candidate context includes outputs, not full per-candidate traces. +- `strategy: "serial"` runs isolated full-program attempts with fresh memory/session IDs. + +## Refine Advice + +`refine(...)` generates advice after a below-threshold round. Advice is appended temporarily to matching `kind: "instruction"` components exposed by `getOptimizableComponents()` and applied through `applyOptimizedComponents()`. + +Rules: + +- Original instruction values are restored in `finally`, on success and error. +- Programs without instruction components continue as best-of-N rounds and mark `adviceApplied: false`. +- Do not add DSPy-style `hint_` signature fields; Ax uses instruction-component advice. + +## Streaming + +Do not use `refine(...)` for streaming. For partial-output safety, use `addStreamingAssert(fieldName, fn, message?)` on `AxGen`. Streaming assertions fail fast within the current attempt with `AxStreamingAssertionError`, then retry with correction feedback when retries remain. diff --git a/.claude/skills/ax-signature/SKILL.md b/.claude/skills/ax-signature/SKILL.md new file mode 100644 index 0000000..8259736 --- /dev/null +++ b/.claude/skills/ax-signature/SKILL.md @@ -0,0 +1,306 @@ +--- +name: ax-signature +description: This skill helps an LLM generate correct DSPy signature code using @ax-llm/ax. Use when the user asks about signatures, s(), f(), field types, string syntax, fluent builder API, validation constraints, or type-safe inputs/outputs. +version: "23.0.0" +--- + +# Ax Signature Reference + +## Signature Syntax + +``` +[description] input1:type, input2:type -> output1:type, output2:type +``` + +## Field Types + +| Type | Syntax | TypeScript | Example | +|------|--------|-----------|---------| +| String | `:string` | `string` | `userName:string` | +| Number | `:number` | `number` | `score:number` | +| Boolean | `:boolean` | `boolean` | `isValid:boolean` | +| JSON | `:json` | `any` | `metadata:json` | +| Date | `:date` | `Date` | `birthDate:date` | +| DateTime | `:datetime` | `Date` | `timestamp:datetime` | +| DateRange | `:dateRange` | `{ start: Date; end: Date }` | `travelDates:dateRange` | +| DateTimeRange | `:datetimeRange` | `{ start: Date; end: Date }` | `meetingWindow:datetimeRange` | +| Image | `:image` | `{mimeType, data}` | `photo:image` (input only) | +| Audio | `:audio` | input: `AxAudioInput`; output: `AxChatAudioOutput` | `recording:audio`, `speech:audio` | +| File | `:file` | `{mimeType, data}` | `document:file` (input only) | +| URL | `:url` | `string` | `website:url` | +| Code | `:code` | `string` | `pythonScript:code` | +| Class | `:class "a, b, c"` | `"a" \| "b" \| "c"` | `mood:class "happy, sad"` | + +Date, datetime, and range fields are AI-friendly but strict. They accept ISO-style values, trim minor whitespace/casing issues, and parse ranges as `{ "start": "...", "end": "..." }`, `[start, end]`, `start/end`, or natural delimiters like `start to end`; invalid values and reversed ranges should fail validation rather than being silently autocorrected. + +## Arrays, Optional, and Internal Fields + +```typescript +'tags:string[] -> processedTags:string[]' // arrays +'query:string, context?:string -> response:string' // optional with ? +'problem:string -> reasoning!:string, solution:string' // internal with ! +``` + +## Four Ways to Create Signatures + +### 1. String-Based (Recommended for simple cases) + +```typescript +import { ax, s } from '@ax-llm/ax'; +const gen = ax('input:string -> output:string'); +const sig = s('query:string -> response:string'); +``` + +### 2. Pure Fluent Builder API + +```typescript +import { f } from '@ax-llm/ax'; +const sig = f() + .input('userMessage', f.string('User input')) + .input('contextData', f.string('Additional context').optional()) + .input('tags', f.string('Keywords').array()) + .output('responseText', f.string('AI response')) + .output('confidenceScore', f.number('Confidence 0-1')) + .output('debugInfo', f.string('Debug info').internal()) + .build(); +``` + +### 3. Standard Schema (zod / valibot / arktype) + +`.input()` and `.output()` accept any [Standard Schema v1](https://standardschema.dev) compatible library — no wrapper, no adapter. Three shapes work everywhere: + +```typescript +import { z } from 'zod'; +import { f } from '@ax-llm/ax'; + +// Shape A: per-field schema — name first, then the schema, then optional ax hints +const sig = f() + .input('contextData', z.string().describe('Background context'), { cache: true }) + .input('userQuestion', z.string().describe('Question to answer')) + .output('reasoning', z.string().describe('Step-by-step thinking'), { internal: true }) + .output('answer', z.string().describe('Final answer')) + .build(); + +// Shape B: whole-object schema — decomposed into fields in declaration order +const sig2 = f() + .description('Answer questions from retrieved context') + .input( + z.object({ + contextData: z.string().describe('Background context'), + userQuestion: z.string().describe('Question to answer'), + }), + { fields: { contextData: { cache: true } } } // companion options map + ) + .output( + z.object({ + reasoning: z.string().describe('Step-by-step thinking'), + answer: z.string().describe('Final answer'), + }), + { fields: { reasoning: { internal: true } } } + ) + .build(); +``` + +Validation constraints from zod flow into ax's prompt validation: + +```typescript +// String constraints: .email(), .url(), .min(), .max(), .regex() +// Number constraints: .min(), .max() +// Arrays: z.array(z.string()) +// Enums: z.enum([...]) — NOTE: enum maps to ax class type, output fields only +const sig3 = f() + .input(z.object({ + emailAddress: z.string().email().describe('Contact email'), + username: z.string().min(3).max(20).describe('Handle'), + score: z.number().min(0).max(100).describe('Numeric score'), + })) + .output(z.object({ + priority: z.enum(['low', 'medium', 'high']).describe('Priority'), + summary: z.string().describe('Result'), + })) + .build(); +``` + +**Companion options** (`AxFieldOptions`) carry ax-specific hints that schema libraries don't represent: + +| Option | Effect | +|--------|--------| +| `{ cache: true }` | Mark input field as a prefix-cache breakpoint | +| `{ internal: true }` | Mark output field as internal scratchpad (stripped from result) | + +The same Standard Schema shapes work on `fn()` tools via `.arg()`, `.returns()`, and `.returnsField()` — argument types are inferred from the schema: + +```typescript +import { z } from 'zod'; +import { fn } from '@ax-llm/ax'; + +// Whole-object zod on a tool — AI-SDK-style +const lookupProduct = fn('lookupProduct') + .description('Look up a product by name and return its current details') + .arg( + z.object({ + productName: z.string().min(1).describe('Exact product name'), + includeSpecs: z.boolean().optional(), + }) + ) + .returns( + z.object({ + price: z.number(), + inStock: z.boolean(), + rating: z.number().min(1).max(5), + }) + ) + .handler(async ({ productName, includeSpecs }) => ({ + price: 79.99, + inStock: true, + rating: 4.3, + })) + .build(); + +// Per-argument form — mix with f.*() args, attach ax hints +const searchDocs = fn('searchDocs') + .description('Search indexed docs') + .arg('query', z.string().min(1), { cache: true }) + .arg('limit', z.number().int().positive().optional()) + .returnsField('results', z.array(z.string())) + .handler(async ({ query }) => []) + .build(); +``` + +### 4. Hybrid + +```typescript +import { s, f } from '@ax-llm/ax'; +const sig = s('base:string -> result:string') + .appendInputField('extra', f.json('Metadata').optional()) + .appendOutputField('score', f.number('Quality score')); +``` + +## Fluent API Reference + +Type creators: +- `f.string(desc)`, `f.number(desc)`, `f.boolean(desc)`, `f.json(desc)` +- `f.image(desc)`, `f.audio(desc)`, `f.file(desc)`, `f.url(desc)` +- `f.email(desc)`, `f.date(desc)`, `f.datetime(desc)`, `f.dateRange(desc)`, `f.datetimeRange(desc)` +- `f.class(['a','b','c'], desc)`, `f.code(desc)` +- `f.object({ field: f.string() }, desc)` + +Chainable modifiers (method chaining only, no nesting): +- `.optional()` - make field optional +- `.array()` / `.array('list description')` - make field an array +- `.internal()` - output only, hidden from final output +- `.cache()` - input only, mark for prompt caching + +```typescript +// Correct: pure fluent chaining +f.string('description').optional().array() +f.string('context').cache().optional() +f.object({ field: f.string() }, 'item desc').array('list desc') + +// Wrong: nested function calls (removed) +f.array(f.string('description')) // REMOVED +f.optional(f.string('description')) // REMOVED +f.internal(f.string('description')) // REMOVED +``` + +## Validation Constraints + +### String Constraints + +```typescript +f.string('username').min(3).max(20) +f.string('email').email() +f.string('website').url() +f.string('birthDate').date() +f.string('timestamp').datetime() +f.string('pattern').regex('^[A-Z0-9]') +``` + +### Number Constraints + +```typescript +f.number('age').min(18).max(120) +f.number('score').min(0).max(100) +``` + +### Complete Validation Example + +```typescript +const sig = f() + .input('formData', f.string('Raw form data')) + .output('user', f.object({ + username: f.string('Username').min(3).max(20), + email: f.string('Email').email(), + age: f.number('Age').min(18).max(120), + bio: f.string('Bio').max(500).optional(), + website: f.string('Website').url().optional(), + tags: f.string('Tag').min(2).max(30).array() + }, 'User profile')) + .build(); +``` + +## Cached Input Fields + +```typescript +const sig = f() + .input('staticContext', f.string('Context').cache()) + .input('userQuery', f.string('Dynamic query')) + .output('answer', f.string('Response')) + .build(); +``` + +## Field Naming Rules + +Good: `userQuestion`, `customerEmail`, `analysisResult`, `confidenceScore` +Bad: `text`, `data`, `input`, `output`, `a`, `x`, `val` (too generic), `1field` (starts with number) + +## Media Type Restrictions + +- Image and file fields are top-level input fields only. +- Audio fields can be top-level inputs or single top-level outputs. +- Audio output fields are scripted speech artifacts: the model returns plain text, then Ax synthesizes `AxChatAudioOutput`. +- Media fields cannot be nested in objects. +- Media arrays are supported for inputs only; output `audio[]` is not supported. + +## Common Patterns + +```typescript +// Chain of Thought +'problem:string -> reasoning!:string, solution:string' + +// Classification +'email:string -> priority:class "urgent, normal, low"' + +// Multi-modal input +'imageData:image, question?:string -> description:string, objects:string[]' + +// Scripted speech output +'question:string -> speech:audio, summary:string' + +// Data Extraction +'invoiceText:string -> invoiceNumber:string, totalAmount:number, lineItems:json[]' + +// With description +'"Answer TypeScript questions" question:string -> answer:string, confidence:number' +``` + +## Critical Rules + +- Use `f()` fluent builder, NOT nested `f.array(f.string())` -- those are removed. +- Field names must be descriptive (not generic like `text`, `data`, `input`). +- Image/file media types are input-only, top-level only; audio may also be a single top-level output. +- `.internal()` / `{ internal: true }` is output-only (for chain-of-thought reasoning). +- `.cache()` / `{ cache: true }` is input-only (for prompt caching). +- Validation errors trigger auto-retry with correction feedback. +- `f.email()`, `f.url()`, `f.date()`, `f.datetime()` are shorthand for `f.string().email()` etc.; `f.dateRange()` and `f.datetimeRange()` return `{ start: Date; end: Date }`. +- `z.enum()` maps to ax's `class` type — only valid on **output** fields. +- For multimodal inputs (images, audio, files) and scripted audio outputs, use `f.image()` / `f.audio()` / `f.file()` — zod has no equivalent. + +## Examples + +Fetch these for full working code: + +- [Standard Schema (zod)](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/standard-schema.ts) — zod with f() and fn(), all three shapes +- [Fluent Signature](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/fluent-signature-example.ts) — native fluent f() API +- [Structured Output](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/structured_output.ts) — structured output with validation +- [Debug Schema](https://raw.githubusercontent.com/ax-llm/ax/refs/heads/main/src/examples/debug_schema.ts) — JSON schema validation From 684df6cbb66d8f4f6c728ef5221746520cc96b14 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 12 Jul 2026 02:15:17 +0000 Subject: [PATCH 7/8] feat: extract ts-autocode-rewrite with AspectJS hot-swappable advice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New workspace package owning both ways a gated candidate becomes real: - guarded source rewriting (applyCandidate, promoteCandidate, revertPromotion, digest) moved from ts-autocode-training with structural RewriteTarget/ RewriteCandidate types, and snapshots now carry the trainable id - an AspectJS Trainable annotation plus @Around aspect weave marked methods so every call dispatches through a hot-swap registry, then one pluggable interceptor, then the original implementation; swapImplementation and restoreImplementation change live behavior without touching source The "use training" directive stays the default marker: the training runtime and register hook now weave discovered methods via annotateTrainable, the @trainable decorator weaves at first construction, and free functions share the same dispatch. training.promote() hot-swaps async targets live and revert() restores both source and implementation. All AspectJS decorators are applied programmatically so consumers stay on standard TC39 decorators. Weaving is covered by aspectjs-style tests (configureTesting(WeaverModule) per test) exercising swap/restore, id scoping, interceptor semantics, idempotent annotation, statics, inheritance, and free functions — entirely in memory, so the tests can never rewrite themselves. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X2LxfjsASWmdVoFANEURWR --- README.md | 10 +- docs/architecture.md | 10 ++ package-lock.json | 42 +++++ package.json | 7 +- packages/rewrite/README.md | 48 ++++++ packages/rewrite/package.json | 45 +++++ packages/rewrite/src/apply.ts | 42 +++++ packages/rewrite/src/aspect.ts | 131 +++++++++++++++ .../{training => rewrite}/src/canonical.ts | 0 packages/rewrite/src/index.ts | 19 +++ packages/rewrite/src/promotion.ts | 57 +++++++ packages/rewrite/test/promotion.test.ts | 68 ++++++++ packages/rewrite/test/weaving.test.ts | 154 ++++++++++++++++++ packages/rewrite/tsconfig.json | 8 + packages/rewrite/tsconfig.test.json | 12 ++ packages/training/package.json | 1 + packages/training/src/engine.ts | 25 +-- packages/training/src/index.ts | 19 ++- packages/training/src/promotion.ts | 50 +----- packages/training/src/source.ts | 2 +- packages/training/src/training.ts | 86 ++++++---- packages/training/test/training.test.ts | 6 +- src/index.ts | 2 + vitest.config.ts | 2 +- 24 files changed, 725 insertions(+), 121 deletions(-) create mode 100644 packages/rewrite/README.md create mode 100644 packages/rewrite/package.json create mode 100644 packages/rewrite/src/apply.ts create mode 100644 packages/rewrite/src/aspect.ts rename packages/{training => rewrite}/src/canonical.ts (100%) create mode 100644 packages/rewrite/src/index.ts create mode 100644 packages/rewrite/src/promotion.ts create mode 100644 packages/rewrite/test/promotion.test.ts create mode 100644 packages/rewrite/test/weaving.test.ts create mode 100644 packages/rewrite/tsconfig.json create mode 100644 packages/rewrite/tsconfig.test.json diff --git a/README.md b/README.md index 8128198..6d7287e 100644 --- a/README.md +++ b/README.md @@ -12,8 +12,9 @@ The normal path keeps the code primitives and agent loop in separate packages: Ax is the default student optimizer. AgentV evaluation and the promotion gate form the teacher. The provider-neutral runtime lives in the independent `ts-autocode-training` package (this package re-exports it with Ax wired in as -the default engine and executor), and iterative coordination is delegated to -the independent `ts-autocode-harness` package. Its single Flue-style callback loop +the default engine and executor), guarded rewriting and hot-swappable AspectJS +interception live in `ts-autocode-rewrite`, and iterative coordination is +delegated to the independent `ts-autocode-harness` package. Its single Flue-style callback loop supports configurable student, teacher, judge, and adversary Deep Agents, MXC execution, and a write-ahead approval bus. Consumers can supply callbacks from their own agent lifecycle or optimization pipeline without coupling it to this @@ -130,6 +131,11 @@ const promoted = await training.promote(run.final.candidate, run.final.decision) await training.revert(promoted.snapshot); ``` +Promotion writes the gated source rewrite and, for async targets, hot-swaps the +running implementation through `ts-autocode-rewrite`'s AspectJS advice — woven +methods dispatch to the promoted candidate immediately, no restart required. +`revert()` restores both the source and the live implementation. + ## Zero-config evolution Load the runtime patch once and directive-marked functions evolve from live diff --git a/docs/architecture.md b/docs/architecture.md index 6ba7a03..650f2e0 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -25,6 +25,16 @@ arguments, synchronous or asynchronous return behavior, and thrown errors. Captured traces use AgentV's `Trace`; spans use official OpenTelemetry and OpenInference APIs. +## Hot-swappable weaving + +`ts-autocode-rewrite` owns candidate application. Marked methods are woven with +an AspectJS `Trainable` annotation whose around advice dispatches through a +hot-swap registry, then a single pluggable interceptor (runtime capture), then +the original implementation. `training.promote()` writes the digest-guarded +source rewrite and swaps async targets live; `revert()` restores both. All +AspectJS decorators are applied programmatically, keeping consumer projects on +standard TC39 decorators. + ## Zero-config runtime patch `ts-autocode/register` installs a `node:module` load hook that appends guarded diff --git a/package-lock.json b/package-lock.json index 33958a4..dcb2fc3 100644 --- a/package-lock.json +++ b/package-lock.json @@ -105,6 +105,24 @@ "integrity": "sha512-4ZeSwiFX3YxB0WSE6x568wM4PVHiYmz3yiOxic6WGKVrE/KIGggMFP/eqUNQhikBKP68IDV0qiILlZAIYnheAQ==", "license": "Apache-2.0" }, + "node_modules/@aspectjs/common": { + "version": "0.5.4", + "resolved": "https://registry.npmjs.org/@aspectjs/common/-/common-0.5.4.tgz", + "integrity": "sha512-tri2wNkswYVfTAqeN++o6q/ylzlzc03yn3xam1awVdxKL8ZmMsDQI9LQ9bVM7v8FhGq3I15RnCmSxrVyl1vUHQ==", + "license": "MIT", + "optionalDependencies": { + "reflect-metadata": "^0.2.2" + } + }, + "node_modules/@aspectjs/core": { + "version": "0.5.4", + "resolved": "https://registry.npmjs.org/@aspectjs/core/-/core-0.5.4.tgz", + "integrity": "sha512-OOSfdXnxSzUvtJ8J3Z3jrmdzc0W8+wDgJIv7TYel9OmP2m+VBFFBjGTKukuD+D5je0s62ySW26a/mbLLQepHFQ==", + "license": "MIT", + "peerDependencies": { + "@aspectjs/common": "^0.5.4" + } + }, "node_modules/@aws-sdk/client-bedrock-runtime": { "version": "3.1084.0", "resolved": "https://registry.npmjs.org/@aws-sdk/client-bedrock-runtime/-/client-bedrock-runtime-3.1084.0.tgz", @@ -3061,6 +3079,13 @@ ], "license": "MIT" }, + "node_modules/reflect-metadata": { + "version": "0.2.2", + "resolved": "https://registry.npmjs.org/reflect-metadata/-/reflect-metadata-0.2.2.tgz", + "integrity": "sha512-urBwgfrvVP/eAyXx4hluJivBKzuEbSQs9rKWCrCkbSxNv8mxPcUZKeuoF3Uy4mJl3Lwprp6yy5/39VWigZ4K6Q==", + "license": "Apache-2.0", + "optional": true + }, "node_modules/retry": { "version": "0.13.1", "resolved": "https://registry.npmjs.org/retry/-/retry-0.13.1.tgz", @@ -3266,6 +3291,10 @@ "resolved": "packages/harness", "link": true }, + "node_modules/ts-autocode-rewrite": { + "resolved": "packages/rewrite", + "link": true + }, "node_modules/ts-autocode-training": { "resolved": "packages/training", "link": true @@ -3570,6 +3599,18 @@ "node": ">=20" } }, + "packages/rewrite": { + "name": "ts-autocode-rewrite", + "version": "0.1.0", + "license": "MIT", + "dependencies": { + "@aspectjs/common": "^0.5.4", + "@aspectjs/core": "^0.5.4" + }, + "engines": { + "node": ">=20" + } + }, "packages/training": { "name": "ts-autocode-training", "version": "0.1.0", @@ -3579,6 +3620,7 @@ "@arizeai/openinference-semantic-conventions": "^2.5.0", "@opentelemetry/api": "^1.9.1", "ts-autocode-harness": "0.1.0", + "ts-autocode-rewrite": "0.1.0", "typescript": "^5.9.3" }, "engines": { diff --git a/package.json b/package.json index 55b5201..4b5a955 100644 --- a/package.json +++ b/package.json @@ -39,14 +39,15 @@ "url": "https://github.com/Tyler-R-Kendrick/ts-autocode/issues" }, "scripts": { - "build": "npm run build:harness && npm run build:training && npm run build:core", + "build": "npm run build:harness && npm run build:rewrite && npm run build:training && npm run build:core", "build:core": "node -e \"require('node:fs').rmSync('dist', { recursive: true, force: true })\" && tsc -p tsconfig.json", "build:harness": "node -e \"require('node:fs').rmSync('packages/harness/dist', { recursive: true, force: true })\" && tsc -p packages/harness/tsconfig.json", - "typecheck": "npm run build:harness && npm run build:training && tsc --noEmit -p tsconfig.test.json && tsc --noEmit -p packages/harness/tsconfig.test.json && tsc --noEmit -p packages/training/tsconfig.test.json", + "typecheck": "npm run build:harness && npm run build:rewrite && npm run build:training && tsc --noEmit -p tsconfig.test.json && tsc --noEmit -p packages/harness/tsconfig.test.json && tsc --noEmit -p packages/rewrite/tsconfig.test.json && tsc --noEmit -p packages/training/tsconfig.test.json", "test": "node test/run.mjs", "check": "npm run typecheck && npm test && npm run build:core", "prepublishOnly": "npm run check", - "build:training": "node -e \"require('node:fs').rmSync('packages/training/dist', { recursive: true, force: true })\" && tsc -p packages/training/tsconfig.json" + "build:training": "node -e \"require('node:fs').rmSync('packages/training/dist', { recursive: true, force: true })\" && tsc -p packages/training/tsconfig.json", + "build:rewrite": "node -e \"require('node:fs').rmSync('packages/rewrite/dist', { recursive: true, force: true })\" && tsc -p packages/rewrite/tsconfig.json" }, "keywords": [ "agentv", diff --git a/packages/rewrite/README.md b/packages/rewrite/README.md new file mode 100644 index 0000000..7cea3b0 --- /dev/null +++ b/packages/rewrite/README.md @@ -0,0 +1,48 @@ +# ts-autocode-rewrite + +Guarded source rewriting and hot-swappable AOP interception for trainable +TypeScript methods. This package owns the two ways a gated candidate becomes +real: + +- **Source rewrite** — `applyCandidate` replaces exactly the discovered method + body behind a digest guard; `promoteCandidate`/`revertPromotion` add + snapshots that refuse to overwrite subsequent edits. +- **Hot-swappable advice** — an [AspectJS](https://www.npmjs.com/package/@aspectjs/core) + `Trainable` annotation and `@Around` aspect weave marked methods so their + live implementation dispatches through a swap registry. + `swapImplementation(id, fn)` replaces behavior in the running process + without touching source; `restoreImplementation(id)` reverts it. + +The `"use training"` literal directive stays the default marker: discovery in +`ts-autocode-training` finds directive-marked methods, and its runtime (or the +`ts-autocode/register` load hook) calls `annotateTrainable(owner, method, id)` +here to weave them. All AspectJS decorators are applied programmatically, so +the package works under both standard and legacy decorator configurations. + +One process-wide `setTrainableInterceptor(fn)` observes every woven +invocation — `ts-autocode-training` wires runtime capture through it. The +interceptor's `proceed()` always resolves the live (possibly swapped) +implementation, so captures reflect what actually ran. + +```ts +import { annotateTrainable, swapImplementation } from "ts-autocode-rewrite"; + +class Router { + route(input: string): string { + "use training"; + return input; + } +} + +annotateTrainable(Router, "route", "Router.route"); +swapImplementation("Router.route", (input) => String(input).toUpperCase()); +new Router().route("abc"); // "ABC" — no source touched +``` + +Most applications should depend on [`ts-autocode`](../../README.md); its +`training.promote()` writes the gated source rewrite **and** hot-swaps async +targets live through this package. + +## License + +[MIT](../../LICENSE) diff --git a/packages/rewrite/package.json b/packages/rewrite/package.json new file mode 100644 index 0000000..3617dcf --- /dev/null +++ b/packages/rewrite/package.json @@ -0,0 +1,45 @@ +{ + "name": "ts-autocode-rewrite", + "version": "0.1.0", + "description": "Guarded source rewriting and hot-swappable AOP interception for trainable TypeScript methods.", + "repository": { + "type": "git", + "url": "https://github.com/Tyler-R-Kendrick/ts-autocode.git", + "directory": "packages/rewrite" + }, + "homepage": "https://github.com/Tyler-R-Kendrick/ts-autocode/tree/main/packages/rewrite#readme", + "bugs": { + "url": "https://github.com/Tyler-R-Kendrick/ts-autocode/issues" + }, + "license": "MIT", + "type": "module", + "sideEffects": false, + "files": [ + "dist" + ], + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + } + }, + "engines": { + "node": ">=20" + }, + "scripts": { + "build": "tsc -p tsconfig.json", + "prepublishOnly": "npm run build" + }, + "dependencies": { + "@aspectjs/common": "^0.5.4", + "@aspectjs/core": "^0.5.4" + }, + "keywords": [ + "aop", + "aspectjs", + "hot-swap", + "rewriter" + ] +} diff --git a/packages/rewrite/src/apply.ts b/packages/rewrite/src/apply.ts new file mode 100644 index 0000000..f8576ea --- /dev/null +++ b/packages/rewrite/src/apply.ts @@ -0,0 +1,42 @@ +import { digest } from "./canonical.js"; + +/** The discovered method-body span a candidate may replace. Structural subset of + * ts-autocode-training's TrainableTarget so either package's targets apply. */ +export interface RewriteTarget { + readonly id: string; + readonly artifactRef: string; + readonly bodyStart: number; + readonly bodyEnd: number; + readonly bodyDigest: string; + readonly indentation: string; +} + +export interface RewriteCandidate { + readonly id: string; + readonly trainableId: string; + readonly target: RewriteTarget; + readonly implementation: string; +} + +/** Replace exactly the discovered method body if it has not changed. */ +export function applyCandidate(source: string, candidate: RewriteCandidate): string { + const { target } = candidate; + if (target.id !== candidate.trainableId) throw new Error("candidate target must match its trainable id"); + const current = source.slice(target.bodyStart, target.bodyEnd); + if (digest(current) !== target.bodyDigest) { + throw new Error(`trainable method changed after optimization started: ${target.id}`); + } + const replacement = formatImplementation(candidate.implementation, target.indentation, source); + return `${source.slice(0, target.bodyStart)}${replacement}${source.slice(target.bodyEnd)}`; +} + +function formatImplementation(implementation: string, methodIndent: string, source: string): string { + const indentUnit = source.includes("\t") ? "\t" : " "; + const bodyIndent = `${methodIndent}${indentUnit}`; + const lines = implementation.split("\n"); + const minimumIndent = Math.min( + ...lines.filter((line) => line.trim()).map((line) => /^\s*/.exec(line)?.[0].length ?? 0), + ); + const normalized = lines.map((line) => `${bodyIndent}${line.slice(Number.isFinite(minimumIndent) ? minimumIndent : 0)}`); + return `\n${normalized.join("\n")}\n${methodIndent}`; +} diff --git a/packages/rewrite/src/aspect.ts b/packages/rewrite/src/aspect.ts new file mode 100644 index 0000000..a0b6902 --- /dev/null +++ b/packages/rewrite/src/aspect.ts @@ -0,0 +1,131 @@ +import { AnnotationFactory, AnnotationKind } from "@aspectjs/common"; +import { Around, Aspect, getWeaver, on, type AroundContext, type JoinPoint } from "@aspectjs/core"; + +/** Marks a method as trainable for the weaver. Applied programmatically by + * `annotateTrainable`, never with decorator syntax, so it works under both + * standard and legacy decorator configurations. */ +export const Trainable = new AnnotationFactory("ts-autocode").create( + AnnotationKind.METHOD, + "Trainable", + // eslint-disable-next-line @typescript-eslint/no-unused-vars + function Trainable(id: string) {}, +); + +export interface TrainableInvocation { + readonly id: string; + readonly methodName: string; + readonly thisValue: unknown; + readonly args: readonly unknown[]; + /** Runs the live implementation: the hot-swapped candidate when one is active, + * otherwise the original joinpoint. */ + readonly proceed: (...args: unknown[]) => unknown; +} + +export type TrainableInterceptor = (invocation: TrainableInvocation) => unknown; + +type AnyMethod = (this: unknown, ...args: unknown[]) => unknown; + +let interceptor: TrainableInterceptor | undefined; +const swaps = new Map(); + +/** One interceptor per process observes every trainable invocation + * (ts-autocode-training wires runtime capture here). */ +export function setTrainableInterceptor(next: TrainableInterceptor | undefined): void { + interceptor = next; +} + +/** Hot-swappable advice: replaces the live implementation for a trainable id. + * Every woven call dispatches through the swap, without touching source. */ +export function swapImplementation(id: string, implementation: AnyMethod): void { + swaps.set(id, implementation); +} + +export function restoreImplementation(id: string): void { + swaps.delete(id); +} + +export function swappedImplementation(id: string): AnyMethod | undefined { + return swaps.get(id); +} + +/** Shared dispatch for the aspect and for wrapped free functions: hot-swap first, + * then the interceptor, then the original implementation. */ +export function dispatchTrainable( + id: string, + methodName: string, + original: AnyMethod, + thisValue: unknown, + args: readonly unknown[], +): unknown { + const proceed = (...next: unknown[]): unknown => { + const active = swaps.get(id) ?? original; + return active.apply(thisValue, next.length > 0 ? next : [...args]); + }; + if (!interceptor) return proceed(); + return interceptor(Object.freeze({ id, methodName, thisValue, args, proceed })); +} + +class TrainableAspectImpl { + intercept(context: AroundContext, joinpoint: JoinPoint, args: unknown[]): unknown { + const found = context.annotations(Trainable).find()[0]; + const id = String(found?.args?.[0] ?? ""); + const methodName = String((context.target as { propertyKey?: unknown }).propertyKey ?? id); + const original: AnyMethod = (...next: unknown[]) => joinpoint(...next); + return dispatchTrainable(id, methodName, original, context.instance, args); + } +} + +let wovenWeaver: unknown; + +/** Idempotent per weaver context; `configureTesting(WeaverModule)` swaps the + * context, after which the next annotate re-enables the aspect. */ +export function enableTrainableWeaving(): void { + const weaver = getWeaver(); + if (weaver === wovenWeaver) return; + wovenWeaver = weaver; + applyLegacyDecorator(Around(on.methods.withAnnotations(Trainable)), TrainableAspectImpl.prototype, "intercept"); + const Enhanced = (Aspect()(TrainableAspectImpl) ?? TrainableAspectImpl) as typeof TrainableAspectImpl; + weaver.enable(new Enhanced()); +} + +const annotatedMethods = new WeakMap>(); + +/** Weaves a class (or static) method for hot-swappable trainable dispatch. + * Walks the prototype chain to the owning container; idempotent per method. */ +export function annotateTrainable( + owner: abstract new (...args: never[]) => unknown, + methodName: string, + id: string, +): void { + const container = owningContainer(owner, methodName); + if (!container) return; + const marked = annotatedMethods.get(container) ?? new Set(); + if (marked.has(methodName)) return; + marked.add(methodName); + annotatedMethods.set(container, marked); + enableTrainableWeaving(); + applyLegacyDecorator(Trainable(id) as LegacyMethodDecorator, container, methodName); +} + +type LegacyMethodDecorator = ( + target: object, + propertyKey: string, + descriptor: PropertyDescriptor, +) => PropertyDescriptor | void; + +function applyLegacyDecorator(decorator: LegacyMethodDecorator, target: object, methodName: string): void { + const descriptor = Object.getOwnPropertyDescriptor(target, methodName); + if (!descriptor) return; + const result = decorator(target, methodName, descriptor); + if (result) Object.defineProperty(target, methodName, result); +} + +function owningContainer(owner: abstract new (...args: never[]) => unknown, methodName: string): object | undefined { + let container: object | null = Object.hasOwn(owner, methodName) ? owner : owner.prototype as object; + while (container && container !== Object.prototype) { + const method = Object.getOwnPropertyDescriptor(container, methodName)?.value as unknown; + if (typeof method === "function") return container; + container = Object.getPrototypeOf(container) as object | null; + } + return undefined; +} diff --git a/packages/training/src/canonical.ts b/packages/rewrite/src/canonical.ts similarity index 100% rename from packages/training/src/canonical.ts rename to packages/rewrite/src/canonical.ts diff --git a/packages/rewrite/src/index.ts b/packages/rewrite/src/index.ts new file mode 100644 index 0000000..58c078b --- /dev/null +++ b/packages/rewrite/src/index.ts @@ -0,0 +1,19 @@ +export { digest, isNonEmptyString } from "./canonical.js"; + +export { applyCandidate } from "./apply.js"; +export type { RewriteCandidate, RewriteTarget } from "./apply.js"; + +export { promoteCandidate, revertPromotion } from "./promotion.js"; +export type { PromotionResult, PromotionSnapshot, RewriteApproval } from "./promotion.js"; + +export { + Trainable, + annotateTrainable, + dispatchTrainable, + enableTrainableWeaving, + restoreImplementation, + setTrainableInterceptor, + swapImplementation, + swappedImplementation, +} from "./aspect.js"; +export type { TrainableInterceptor, TrainableInvocation } from "./aspect.js"; diff --git a/packages/rewrite/src/promotion.ts b/packages/rewrite/src/promotion.ts new file mode 100644 index 0000000..54aee79 --- /dev/null +++ b/packages/rewrite/src/promotion.ts @@ -0,0 +1,57 @@ +import { applyCandidate, type RewriteCandidate } from "./apply.js"; + +/** Structural subset of ts-autocode-training's PromotionDecision. */ +export interface RewriteApproval { + readonly candidateId: string; + readonly promote: boolean; +} + +export interface PromotionSnapshot { + readonly candidateId: string; + readonly trainableId: string; + readonly artifactRef: string; + readonly startOffset: number; + readonly previous: string; + readonly promoted: string; +} + +export interface PromotionResult { + readonly source: string; + readonly snapshot: PromotionSnapshot; +} + +export function promoteCandidate({ + source, + candidate, + decision, +}: { + source: string; + candidate: RewriteCandidate; + decision: RewriteApproval; +}): PromotionResult { + if (!decision.promote || decision.candidateId !== candidate.id) { + throw new Error("candidate has not passed the promotion gate"); + } + const updated = applyCandidate(source, candidate); + const previous = source.slice(candidate.target.bodyStart, candidate.target.bodyEnd); + const promotedLength = updated.length - source.length + previous.length; + return Object.freeze({ + source: updated, + snapshot: Object.freeze({ + candidateId: candidate.id, + trainableId: candidate.trainableId, + artifactRef: candidate.target.artifactRef, + startOffset: candidate.target.bodyStart, + previous, + promoted: updated.slice(candidate.target.bodyStart, candidate.target.bodyStart + promotedLength), + }), + }); +} + +export function revertPromotion(source: string, snapshot: PromotionSnapshot): string { + const endOffset = snapshot.startOffset + snapshot.promoted.length; + if (source.slice(snapshot.startOffset, endOffset) !== snapshot.promoted) { + throw new Error("promoted method changed before revert"); + } + return `${source.slice(0, snapshot.startOffset)}${snapshot.previous}${source.slice(endOffset)}`; +} diff --git a/packages/rewrite/test/promotion.test.ts b/packages/rewrite/test/promotion.test.ts new file mode 100644 index 0000000..45e06e9 --- /dev/null +++ b/packages/rewrite/test/promotion.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, it } from "vitest"; + +import { + applyCandidate, + digest, + promoteCandidate, + revertPromotion, + type RewriteCandidate, + type RewriteTarget, +} from "../src/index.js"; + +const source = `class Router { + route(input: string): string { + "use training"; + return input; + } +}`; + +function targetFor(text: string): RewriteTarget { + const bodyStart = text.indexOf('"use training";') + '"use training";'.length; + const bodyEnd = text.lastIndexOf("\n }"); + return { + id: "Router.route", + artifactRef: "memory://router.ts", + bodyStart, + bodyEnd, + bodyDigest: digest(text.slice(bodyStart, bodyEnd)), + indentation: " ", + }; +} + +function candidateFor(text: string, implementation: string): RewriteCandidate { + return { id: "candidate-1", trainableId: "Router.route", target: targetFor(text), implementation }; +} + +describe("guarded source rewrite", () => { + it("replaces exactly the discovered body and preserves the directive", () => { + const updated = applyCandidate(source, candidateFor(source, "return input.toUpperCase();")); + expect(updated).toContain('"use training";'); + expect(updated).toContain("return input.toUpperCase();"); + expect(updated).not.toContain("return input;\n }"); + }); + + it("refuses stale targets whose body changed after discovery", () => { + const candidate = candidateFor(source, "return input.toUpperCase();"); + const drifted = source.replace("return input;", "return input.trim();"); + expect(() => applyCandidate(drifted, candidate)).toThrow("changed after optimization started"); + }); + + it("promotes only gate-approved candidates and records a revertible snapshot", () => { + const candidate = candidateFor(source, "return input.toUpperCase();"); + expect(() => promoteCandidate({ source, candidate, decision: { candidateId: candidate.id, promote: false } })) + .toThrow("has not passed the promotion gate"); + expect(() => promoteCandidate({ source, candidate, decision: { candidateId: "other", promote: true } })) + .toThrow("has not passed the promotion gate"); + + const promoted = promoteCandidate({ source, candidate, decision: { candidateId: candidate.id, promote: true } }); + expect(promoted.snapshot.trainableId).toBe("Router.route"); + expect(revertPromotion(promoted.source, promoted.snapshot)).toBe(source); + }); + + it("refuses to revert over subsequent edits", () => { + const candidate = candidateFor(source, "return input.toUpperCase();"); + const promoted = promoteCandidate({ source, candidate, decision: { candidateId: candidate.id, promote: true } }); + const edited = promoted.source.replace("toUpperCase", "toLowerCase"); + expect(() => revertPromotion(edited, promoted.snapshot)).toThrow("changed before revert"); + }); +}); diff --git a/packages/rewrite/test/weaving.test.ts b/packages/rewrite/test/weaving.test.ts new file mode 100644 index 0000000..f690f00 --- /dev/null +++ b/packages/rewrite/test/weaving.test.ts @@ -0,0 +1,154 @@ +import { configureTesting } from "@aspectjs/common/testing"; +import { WeaverModule } from "@aspectjs/core"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { + annotateTrainable, + dispatchTrainable, + restoreImplementation, + setTrainableInterceptor, + swapImplementation, + swappedImplementation, + type TrainableInvocation, +} from "../src/index.js"; + +// Weaving is exercised entirely in memory: hot-swapped advice, never source +// edits, so these tests can never rewrite themselves. +describe("trainable weaving", () => { + beforeEach(() => { + configureTesting(WeaverModule); + setTrainableInterceptor(undefined); + restoreImplementation("Router.route"); + restoreImplementation("Router.fallback"); + restoreImplementation("Static.echo"); + restoreImplementation("free.normalize"); + }); + + it("weaves annotated methods and leaves sibling methods untouched", () => { + class Router { + route(input: string): string { return input; } + fallback(input: string): string { return input; } + } + annotateTrainable(Router, "route", "Router.route"); + const seen: string[] = []; + setTrainableInterceptor((invocation) => { + seen.push(invocation.id); + return invocation.proceed(); + }); + + const router = new Router(); + expect(router.route("abc")).toBe("abc"); + expect(router.fallback("abc")).toBe("abc"); + expect(seen).toEqual(["Router.route"]); + }); + + it("hot-swaps the live implementation and restores the original", () => { + class Router { + route(input: string): string { return input; } + } + annotateTrainable(Router, "route", "Router.route"); + const router = new Router(); + + expect(router.route("abc")).toBe("abc"); + swapImplementation("Router.route", (input) => String(input).toUpperCase()); + expect(router.route("abc")).toBe("ABC"); + expect(typeof swappedImplementation("Router.route")).toBe("function"); + restoreImplementation("Router.route"); + expect(router.route("abc")).toBe("abc"); + expect(swappedImplementation("Router.route")).toBeUndefined(); + }); + + it("applies swaps through the interceptor's proceed and preserves this/args", () => { + class Router { + prefix = "id:"; + route(input: string): string { return `${this.prefix}${input}`; } + } + annotateTrainable(Router, "route", "Router.route"); + const invocations: TrainableInvocation[] = []; + setTrainableInterceptor((invocation) => { + invocations.push(invocation); + return invocation.proceed(); + }); + swapImplementation("Router.route", function (this: unknown, input) { + return `${(this as Router).prefix}${String(input).toUpperCase()}`; + }); + + expect(new Router().route("abc")).toBe("id:ABC"); + expect(invocations[0]?.id).toBe("Router.route"); + expect(invocations[0]?.methodName).toBe("route"); + expect(invocations[0]?.args).toEqual(["abc"]); + }); + + it("keeps swaps scoped to their trainable id", () => { + class Router { + route(input: string): string { return input; } + fallback(input: string): string { return input; } + } + annotateTrainable(Router, "route", "Router.route"); + annotateTrainable(Router, "fallback", "Router.fallback"); + swapImplementation("Router.route", () => "swapped"); + + const router = new Router(); + expect(router.route("abc")).toBe("swapped"); + expect(router.fallback("abc")).toBe("abc"); + }); + + it("annotates idempotently so advice runs once per call", () => { + class Router { + route(input: string): string { return input; } + } + annotateTrainable(Router, "route", "Router.route"); + annotateTrainable(Router, "route", "Router.route"); + const interceptor = vi.fn((invocation: TrainableInvocation) => invocation.proceed()); + setTrainableInterceptor(interceptor); + + expect(new Router().route("abc")).toBe("abc"); + expect(interceptor).toHaveBeenCalledTimes(1); + }); + + it("weaves static methods and inherited methods through the owning container", () => { + class Static { + static echo(input: string): string { return input; } + } + annotateTrainable(Static, "echo", "Static.echo"); + swapImplementation("Static.echo", (input) => `static:${String(input)}`); + expect(Static.echo("x")).toBe("static:x"); + + class Base { + route(input: string): string { return input; } + } + class Derived extends Base {} + annotateTrainable(Derived, "route", "Router.route"); + swapImplementation("Router.route", () => "woven"); + expect(new Derived().route("abc")).toBe("woven"); + expect(new Base().route("abc")).toBe("woven"); + }); + + it("dispatches wrapped free functions through the same swap registry", () => { + const normalize = (input: string): string => input.trim(); + const call = (input: string): unknown => + dispatchTrainable("free.normalize", "normalize", normalize as (...args: unknown[]) => unknown, undefined, [input]); + + expect(call(" x ")).toBe("x"); + swapImplementation("free.normalize", (input) => String(input).trim().toUpperCase()); + expect(call(" x ")).toBe("X"); + restoreImplementation("free.normalize"); + expect(call(" x ")).toBe("x"); + }); + + it("interceptors can observe without altering results, and clearing them restores plain dispatch", () => { + class Router { + route(input: string): string { return input; } + } + annotateTrainable(Router, "route", "Router.route"); + const router = new Router(); + const interceptor = vi.fn((invocation: TrainableInvocation) => invocation.proceed()); + setTrainableInterceptor(interceptor); + expect(router.route("abc")).toBe("abc"); + expect(interceptor).toHaveBeenCalledTimes(1); + + setTrainableInterceptor(undefined); + expect(router.route("abc")).toBe("abc"); + expect(interceptor).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/rewrite/tsconfig.json b/packages/rewrite/tsconfig.json new file mode 100644 index 0000000..cdc3a24 --- /dev/null +++ b/packages/rewrite/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.json", + "compilerOptions": { + "rootDir": "src", + "outDir": "dist" + }, + "include": ["src"] +} diff --git a/packages/rewrite/tsconfig.test.json b/packages/rewrite/tsconfig.test.json new file mode 100644 index 0000000..250b99d --- /dev/null +++ b/packages/rewrite/tsconfig.test.json @@ -0,0 +1,12 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "noEmit": true, + "rootDir": ".", + "outDir": null, + "declaration": false, + "declarationMap": false, + "sourceMap": false + }, + "include": ["src", "test"] +} diff --git a/packages/training/package.json b/packages/training/package.json index 0c2de03..4a121c7 100644 --- a/packages/training/package.json +++ b/packages/training/package.json @@ -37,6 +37,7 @@ "@arizeai/openinference-semantic-conventions": "^2.5.0", "@opentelemetry/api": "^1.9.1", "ts-autocode-harness": "0.1.0", + "ts-autocode-rewrite": "0.1.0", "typescript": "^5.9.3" }, "keywords": [ diff --git a/packages/training/src/engine.ts b/packages/training/src/engine.ts index 1dd2739..1ac3feb 100644 --- a/packages/training/src/engine.ts +++ b/packages/training/src/engine.ts @@ -1,7 +1,7 @@ import type { EvaluationResult, EvalTestInput } from "@agentv/core"; +import { digest } from "ts-autocode-rewrite"; import ts from "typescript"; -import { digest } from "./canonical.js"; import type { TrainingRecord } from "./records.js"; import type { TrainableTarget } from "./source.js"; import type { TrainableId } from "./token.js"; @@ -80,18 +80,6 @@ export async function optimizeCandidate( return Object.freeze(structuredClone(candidate)); } -/** Replace exactly the discovered method body if it has not changed. */ -export function applyCandidate(source: string, candidate: CandidatePatch): string { - const { target } = candidate; - if (target.id !== candidate.trainableId) throw new Error("candidate target must match its trainable id"); - const current = source.slice(target.bodyStart, target.bodyEnd); - if (digest(current) !== target.bodyDigest) { - throw new Error(`trainable method changed after optimization started: ${target.id}`); - } - const replacement = formatImplementation(candidate.implementation, target.indentation, source); - return `${source.slice(0, target.bodyStart)}${replacement}${source.slice(target.bodyEnd)}`; -} - function validateRequest(request: OptimizeRequest): void { if (!request.objective.trim()) throw new TypeError("optimization objective must be a non-empty string"); if (request.target.id !== request.trainableId) throw new Error("trainable target must match the request id"); @@ -120,14 +108,3 @@ function validateImplementation(target: TrainableTarget, implementation: string) throw new SyntaxError(`engine returned invalid TypeScript for ${target.id}`); } } - -function formatImplementation(implementation: string, methodIndent: string, source: string): string { - const indentUnit = source.includes("\t") ? "\t" : " "; - const bodyIndent = `${methodIndent}${indentUnit}`; - const lines = implementation.split("\n"); - const minimumIndent = Math.min( - ...lines.filter((line) => line.trim()).map((line) => /^\s*/.exec(line)?.[0].length ?? 0), - ); - const normalized = lines.map((line) => `${bodyIndent}${line.slice(Number.isFinite(minimumIndent) ? minimumIndent : 0)}`); - return `\n${normalized.join("\n")}\n${methodIndent}`; -} diff --git a/packages/training/src/index.ts b/packages/training/src/index.ts index edba4bf..53bd741 100644 --- a/packages/training/src/index.ts +++ b/packages/training/src/index.ts @@ -28,7 +28,6 @@ export type { TrainableId, TrainableIdentity, TrainableToken } from "./token.js" export { discoverInSource, discoverTrainables } from "./source.js"; export type { SourceSettings, TrainableTarget } from "./source.js"; -export { applyCandidate } from "./engine.js"; export type { BoundEvaluation, CandidatePatch, @@ -40,13 +39,17 @@ export type { TrainingEngine, } from "./engine.js"; -export { evaluatePromotionGate, promoteCandidate, revertPromotion } from "./promotion.js"; -export type { - PromotionDecision, - PromotionGateInput, - PromotionResult, - PromotionSnapshot, -} from "./promotion.js"; +export { evaluatePromotionGate } from "./promotion.js"; +export type { PromotionDecision, PromotionGateInput } from "./promotion.js"; + +export { + applyCandidate, + promoteCandidate, + restoreImplementation, + revertPromotion, + swapImplementation, +} from "ts-autocode-rewrite"; +export type { PromotionResult, PromotionSnapshot } from "ts-autocode-rewrite"; export { createMemoryTrainingStore } from "./records.js"; export type { TrainingRecord, TrainingStore } from "./records.js"; diff --git a/packages/training/src/promotion.ts b/packages/training/src/promotion.ts index e37682d..3a9e667 100644 --- a/packages/training/src/promotion.ts +++ b/packages/training/src/promotion.ts @@ -1,6 +1,6 @@ import type { EvaluationResult } from "@agentv/core"; -import { applyCandidate, type BoundEvaluation, type CandidatePatch } from "./engine.js"; +import type { BoundEvaluation, CandidatePatch } from "./engine.js"; export interface PromotionGateInput { readonly candidate: CandidatePatch; @@ -19,19 +19,6 @@ export interface PromotionDecision { readonly passRate: number; } -export interface PromotionSnapshot { - readonly candidateId: string; - readonly artifactRef: string; - readonly startOffset: number; - readonly previous: string; - readonly promoted: string; -} - -export interface PromotionResult { - readonly source: string; - readonly snapshot: PromotionSnapshot; -} - export async function evaluatePromotionGate(input: PromotionGateInput): Promise { const minScore = input.minScore ?? 0.8; const minPassRate = input.minPassRate ?? 1; @@ -62,41 +49,6 @@ export async function evaluatePromotionGate(input: PromotionGateInput): Promise< return Object.freeze({ candidateId: input.candidate.id, promote: failures.length === 0, failures, meanScore, passRate }); } -export function promoteCandidate({ - source, - candidate, - decision, -}: { - source: string; - candidate: CandidatePatch; - decision: PromotionDecision; -}): PromotionResult { - if (!decision.promote || decision.candidateId !== candidate.id) { - throw new Error("candidate has not passed the promotion gate"); - } - const updated = applyCandidate(source, candidate); - const previous = source.slice(candidate.target.bodyStart, candidate.target.bodyEnd); - const promotedLength = updated.length - source.length + previous.length; - return Object.freeze({ - source: updated, - snapshot: Object.freeze({ - candidateId: candidate.id, - artifactRef: candidate.target.artifactRef, - startOffset: candidate.target.bodyStart, - previous, - promoted: updated.slice(candidate.target.bodyStart, candidate.target.bodyStart + promotedLength), - }), - }); -} - -export function revertPromotion(source: string, snapshot: PromotionSnapshot): string { - const endOffset = snapshot.startOffset + snapshot.promoted.length; - if (source.slice(snapshot.startOffset, endOffset) !== snapshot.promoted) { - throw new Error("promoted method changed before revert"); - } - return `${source.slice(0, snapshot.startOffset)}${snapshot.previous}${source.slice(endOffset)}`; -} - function passed(result: EvaluationResult, threshold: number): boolean { return result.executionStatus !== "execution_error" && result.score >= threshold; } diff --git a/packages/training/src/source.ts b/packages/training/src/source.ts index 4697b86..89f269c 100644 --- a/packages/training/src/source.ts +++ b/packages/training/src/source.ts @@ -2,7 +2,7 @@ import { dirname, extname, resolve } from "node:path"; import ts from "typescript"; -import { digest } from "./canonical.js"; +import { digest } from "ts-autocode-rewrite"; import { trainableIdFromKey, type TrainableId } from "./token.js"; const trainingDirective = "use training"; diff --git a/packages/training/src/training.ts b/packages/training/src/training.ts index 6182676..6a56b88 100644 --- a/packages/training/src/training.ts +++ b/packages/training/src/training.ts @@ -6,6 +6,17 @@ import { buildTraceFromMessages, getTextContent, type EvalConfig, type EvalTestI import { OpenInferenceSpanKind, SemanticConventions } from "@arizeai/openinference-semantic-conventions"; import { SpanStatusCode, trace, type Attributes, type Span, type Tracer } from "@opentelemetry/api"; import { defineTrainingHarness, WriteAheadAgentBus, type JudgeRequest } from "ts-autocode-harness"; +import { + annotateTrainable, + dispatchTrainable, + promoteCandidate, + restoreImplementation, + revertPromotion, + setTrainableInterceptor, + swapImplementation, + type PromotionResult, + type PromotionSnapshot, +} from "ts-autocode-rewrite"; import { optimizeCandidate, @@ -16,14 +27,7 @@ import { type TrainingEngine, } from "./engine.js"; import { evaluateTrainable, type TrainableEvalRun } from "./evaluation.js"; -import { - evaluatePromotionGate, - promoteCandidate, - revertPromotion, - type PromotionDecision, - type PromotionResult, - type PromotionSnapshot, -} from "./promotion.js"; +import { evaluatePromotionGate, type PromotionDecision } from "./promotion.js"; import { createMemoryTrainingStore, type TrainingRecord, type TrainingStore } from "./records.js"; import { discoverTrainables, @@ -441,12 +445,25 @@ class TrainingRuntime implements Training { const source = await readFile(candidate.target.artifactRef, "utf8"); const promoted = promoteCandidate({ source, candidate, decision }); await writeFile(candidate.target.artifactRef, promoted.source, "utf8"); + this.#hotSwap(candidate); return promoted; } async revert(snapshot: PromotionSnapshot): Promise { const source = await readFile(snapshot.artifactRef, "utf8"); await writeFile(snapshot.artifactRef, revertPromotion(source, snapshot), "utf8"); + restoreImplementation(snapshot.trainableId); + } + + /** Promoted candidates go live in-process through the hot-swappable advice. + * Only async targets swap: the executor returns a promise, so swapping a + * synchronous method would change its calling convention. */ + #hotSwap(candidate: CandidatePatch): void { + if (!candidate.target.async) return; + const executor = this.#settings.executor ?? defaultProviders.executor; + if (!executor) return; + swapImplementation(candidate.trainableId, (...args: unknown[]) => + executor(candidate.target, candidate.implementation, args)); } async flush(): Promise { @@ -614,9 +631,22 @@ export const training: Training = Object.freeze({ const wrappedMarker = Symbol.for("ts-autocode.wrapped"); +// Every woven or wrapped trainable dispatches through ts-autocode-rewrite's +// hot-swappable advice; this interceptor routes each invocation into runtime +// capture, and `proceed` resolves the live (possibly swapped) implementation. +setTrainableInterceptor(({ id, methodName, thisValue, args, proceed }) => + runtime().invoke( + thisValue, + function (this: unknown, ...next: unknown[]) { return proceed(...next); }, + [...args], + defineTrainable(id), + methodName, + )); + /** Decorator form: `@trainable()`. Identity is inferred from the decorated class and * method; pass a symbol (for example `defineTrainable("Router.route").symbol`) only - * to override the inferred id. */ + * to override the inferred id. The method is woven through the ts-autocode-rewrite + * aspect at first construction, so promoted candidates can hot-swap it. */ export function trainable(identity?: symbol): TrainableDecorator { if (identity !== undefined && typeof identity !== "symbol") { throw new TypeError("trainable identity must be a symbol; omit it to infer from the decorated method"); @@ -627,45 +657,37 @@ export function trainable(identity?: symbol): TrainableDecorator { context: ClassMethodDecoratorContext Result>, ) { const name = String(context.name); - let token = explicit; - const wrapped = function (this: This, ...args: Args): Result { - token ??= defineTrainable(`${inferredClassName(this) ?? "Anonymous"}.${name}`); - return runtime().invoke(this, method, args, token, name); - }; - return markWrapped(wrapped); + context.addInitializer(function (this: This) { + const owner = (context.static ? this : (this as object).constructor) as abstract new (...args: never[]) => unknown; + const id = explicit?.id ?? `${inferredClassName(context.static ? owner : this) ?? "Anonymous"}.${name}`; + annotateTrainable(owner, name, id); + }); + return method; }; } -/** Load-time instrumentation (`ts-autocode/register`): capture-wrap a directive-marked - * function. Idempotent — already-wrapped functions (decorator or register) pass through. */ +/** Load-time instrumentation (`ts-autocode/register`): wrap a directive-marked free + * function through the same hot-swappable dispatch as woven methods. Idempotent. */ export function wrapTrainable unknown>(fn: F, id: string): F { if ((fn as Partial>)[wrappedMarker]) return fn; - const token = defineTrainable(id); - const name = fn.name || token.id; + const name = fn.name || id; const method = fn as unknown as (this: unknown, ...args: unknown[]) => unknown; const wrapped = function (this: unknown, ...args: unknown[]): unknown { - return runtime().invoke(this, method, args, token, name); + return dispatchTrainable(id, name, method, this, args); }; Object.defineProperty(wrapped, "name", { value: name, configurable: true }); - return markWrapped(wrapped) as unknown as F; + Object.defineProperty(wrapped, wrappedMarker, { value: true }); + return wrapped as unknown as F; } -/** Load-time instrumentation (`ts-autocode/register`): capture-wrap a directive-marked - * class method in place. Idempotent and tolerant of missing members. */ +/** Load-time instrumentation (`ts-autocode/register`): weave a directive-marked + * class method through the ts-autocode-rewrite aspect. Idempotent. */ export function instrumentTrainable( owner: abstract new (...args: never[]) => unknown, methodName: string, id: string, ): void { - const container = (Object.hasOwn(owner, methodName) ? owner : owner.prototype) as Record; - const method = container?.[methodName]; - if (typeof method !== "function") return; - container[methodName] = wrapTrainable(method as (...args: never[]) => unknown, id); -} - -function markWrapped(fn: F): F { - Object.defineProperty(fn, wrappedMarker, { value: true }); - return fn; + annotateTrainable(owner, methodName, id); } function inferredClassName(thisValue: unknown): string | undefined { diff --git a/packages/training/test/training.test.ts b/packages/training/test/training.test.ts index 71f3ee4..8e0aa16 100644 --- a/packages/training/test/training.test.ts +++ b/packages/training/test/training.test.ts @@ -285,6 +285,7 @@ function applyMethodDecorator o ): void { const prototype = constructor.prototype as Record; const method = prototype[name] as (...args: unknown[]) => unknown; + const initializers: Array<(this: object) => void> = []; const replacement = decorator(method, { kind: "method", name, @@ -294,8 +295,11 @@ function applyMethodDecorator o has: (value: unknown) => name in (value as object), get: (value: unknown) => (value as Record)[name] as (...args: unknown[]) => unknown, }, - addInitializer() {}, + addInitializer(initializer: (this: object) => void) { + initializers.push(initializer); + }, metadata: undefined, } as unknown as ClassMethodDecoratorContext); Object.defineProperty(prototype, name, { value: replacement, configurable: true, writable: true }); + for (const initializer of initializers) initializer.call(Object.create(constructor.prototype) as object); } diff --git a/src/index.ts b/src/index.ts index 9e00965..f872c05 100644 --- a/src/index.ts +++ b/src/index.ts @@ -13,7 +13,9 @@ export { discoverTrainables, evaluatePromotionGate, promoteCandidate, + restoreImplementation, revertPromotion, + swapImplementation, trainable, training, } from "ts-autocode-training"; diff --git a/vitest.config.ts b/vitest.config.ts index 0bda23a..d6137d9 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -2,6 +2,6 @@ import { defineConfig } from "vitest/config"; export default defineConfig({ test: { - include: ["test/**/*.test.ts", "packages/harness/test/**/*.test.ts", "packages/training/test/**/*.test.ts"], + include: ["test/**/*.test.ts", "packages/harness/test/**/*.test.ts", "packages/rewrite/test/**/*.test.ts", "packages/training/test/**/*.test.ts"], }, }); From c7710809a09ce827c0ef18f1c36fb219a01f3e42 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 12 Jul 2026 02:31:33 +0000 Subject: [PATCH 8/8] refactor!: make ts-autocode-rewrite a generic marker-configured engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The rewrite package no longer knows about training. A consumer registers a "use " marker once via the single configureRewrite entry point, and marking a method with that directive is all that's needed after that — weaving and hot-swapping happen through the configured behavior, not through explicit consumer calls. - rename Trainable annotation -> Rewrite (carries id + marker); annotateTrainable -> annotateRewrite; setTrainableInterceptor -> per-marker configureRewrite; dispatchTrainable -> dispatchRewrite - "use " is shorthand for annotateRewrite; markers are normalized and used as annotation configuration keys, so multiple rewrite behaviors coexist - swapImplementation/annotateRewrite remain exported for tests and advanced orchestration but are off the default consumer path - ts-autocode-training registers "use training" with its capture interceptor; the @trainable decorator and register hook weave under that marker implicitly - infer @trainable identity from the declaring class (not a runtime subclass) - hot-swap wrapper is a normal function forwarding the receiver to executors - formatImplementation picks the indent unit from the method's own indentation Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X2LxfjsASWmdVoFANEURWR --- packages/rewrite/README.md | 75 +++++++++++------ packages/rewrite/src/apply.ts | 4 +- packages/rewrite/src/aspect.ts | 116 +++++++++++++++++++------- packages/rewrite/src/index.ts | 16 ++-- packages/rewrite/test/weaving.test.ts | 78 ++++++++--------- packages/training/src/engine.ts | 6 +- packages/training/src/training.ts | 76 ++++++++++------- 7 files changed, 238 insertions(+), 133 deletions(-) diff --git a/packages/rewrite/README.md b/packages/rewrite/README.md index 7cea3b0..e6c946c 100644 --- a/packages/rewrite/README.md +++ b/packages/rewrite/README.md @@ -1,47 +1,68 @@ # ts-autocode-rewrite -Guarded source rewriting and hot-swappable AOP interception for trainable -TypeScript methods. This package owns the two ways a gated candidate becomes -real: +Guarded source rewriting and hot-swappable AOP interception, driven by a +configurable `"use "` marker. The package is general — it knows nothing +about training. A consumer registers a marker and its behavior once, and +marking a method with that directive is all that's needed after that. + +## Two ways a candidate becomes real - **Source rewrite** — `applyCandidate` replaces exactly the discovered method body behind a digest guard; `promoteCandidate`/`revertPromotion` add snapshots that refuse to overwrite subsequent edits. - **Hot-swappable advice** — an [AspectJS](https://www.npmjs.com/package/@aspectjs/core) - `Trainable` annotation and `@Around` aspect weave marked methods so their - live implementation dispatches through a swap registry. - `swapImplementation(id, fn)` replaces behavior in the running process - without touching source; `restoreImplementation(id)` reverts it. - -The `"use training"` literal directive stays the default marker: discovery in -`ts-autocode-training` finds directive-marked methods, and its runtime (or the -`ts-autocode/register` load hook) calls `annotateTrainable(owner, method, id)` -here to weave them. All AspectJS decorators are applied programmatically, so -the package works under both standard and legacy decorator configurations. - -One process-wide `setTrainableInterceptor(fn)` observes every woven -invocation — `ts-autocode-training` wires runtime capture through it. The -interceptor's `proceed()` always resolves the live (possibly swapped) -implementation, so captures reflect what actually ran. + `Rewrite` annotation and `@Around` aspect weave marked methods so their live + implementation dispatches through a swap registry. Promotion swaps behavior + in the running process without touching source. + +## Configure a marker, then just mark methods + +`configureRewrite` is the single entry point. It binds a `"use "` marker +to its rewrite behavior (an optional per-invocation interceptor). After that, +the `"use "` directive is the shorthand — a consumer's discovery or load +hook weaves each marked method, and promotion drives the swap: ```ts -import { annotateTrainable, swapImplementation } from "ts-autocode-rewrite"; +import { configureRewrite } from "ts-autocode-rewrite"; +// A consumer registers its marker once (ts-autocode-training registers +// "use training" with its runtime-capture interceptor). +configureRewrite({ + marker: "use audit", + intercept: (call) => { + log(call.id, call.args); + return call.proceed(); // resolves the live (possibly swapped) implementation + }, +}); +``` + +```ts class Router { route(input: string): string { - "use training"; + "use audit"; // the marker is the only thing a consumer writes return input; } } - -annotateTrainable(Router, "route", "Router.route"); -swapImplementation("Router.route", (input) => String(input).toUpperCase()); -new Router().route("abc"); // "ABC" — no source touched ``` -Most applications should depend on [`ts-autocode`](../../README.md); its -`training.promote()` writes the gated source rewrite **and** hot-swaps async -targets live through this package. +Different markers route to different configurations, so several rewrite +behaviors can coexist in one process. The `"use "` marker is normalized +and used as the annotation's configuration key. + +## Advanced / test helpers + +`annotateRewrite(owner, method, id, marker)` weaves a method directly, and +`swapImplementation(id, fn)` / `restoreImplementation(id)` change live behavior +without touching source. These back the shorthand above and are exported for +tests and custom orchestration — they are not part of the normal consumer path, +which is: configure a marker, mark methods, promote. + +Every AspectJS decorator is applied programmatically, so the package works +under both standard and legacy decorator configurations. + +Most applications should depend on [`ts-autocode`](../../README.md), which +configures the `"use training"` marker and drives promotion (source rewrite plus +live hot-swap of async targets) for you. ## License diff --git a/packages/rewrite/src/apply.ts b/packages/rewrite/src/apply.ts index f8576ea..0087e8f 100644 --- a/packages/rewrite/src/apply.ts +++ b/packages/rewrite/src/apply.ts @@ -31,7 +31,9 @@ export function applyCandidate(source: string, candidate: RewriteCandidate): str } function formatImplementation(implementation: string, methodIndent: string, source: string): string { - const indentUnit = source.includes("\t") ? "\t" : " "; + // Match the method's own indentation style rather than the whole file's, + // so a tab-indented method in a mostly-spaces file still gets tabs. + const indentUnit = methodIndent.includes("\t") ? "\t" : source.includes("\t") ? "\t" : " "; const bodyIndent = `${methodIndent}${indentUnit}`; const lines = implementation.split("\n"); const minimumIndent = Math.min( diff --git a/packages/rewrite/src/aspect.ts b/packages/rewrite/src/aspect.ts index a0b6902..488cee9 100644 --- a/packages/rewrite/src/aspect.ts +++ b/packages/rewrite/src/aspect.ts @@ -1,18 +1,20 @@ import { AnnotationFactory, AnnotationKind } from "@aspectjs/common"; import { Around, Aspect, getWeaver, on, type AroundContext, type JoinPoint } from "@aspectjs/core"; -/** Marks a method as trainable for the weaver. Applied programmatically by - * `annotateTrainable`, never with decorator syntax, so it works under both - * standard and legacy decorator configurations. */ -export const Trainable = new AnnotationFactory("ts-autocode").create( +/** Marks a method for the weaver. Carries the stable id and the configured + * marker (e.g. `"use training"`) so dispatch can look up that marker's config. + * Applied programmatically by `annotateRewrite`, never with decorator syntax, + * so it works under both standard and legacy decorator configurations. */ +export const Rewrite = new AnnotationFactory("ts-autocode").create( AnnotationKind.METHOD, - "Trainable", + "Rewrite", // eslint-disable-next-line @typescript-eslint/no-unused-vars - function Trainable(id: string) {}, + function Rewrite(id: string, marker: string) {}, ); -export interface TrainableInvocation { +export interface RewriteInvocation { readonly id: string; + readonly marker: string; readonly methodName: string; readonly thisValue: unknown; readonly args: readonly unknown[]; @@ -21,21 +23,53 @@ export interface TrainableInvocation { readonly proceed: (...args: unknown[]) => unknown; } -export type TrainableInterceptor = (invocation: TrainableInvocation) => unknown; +export type RewriteInterceptor = (invocation: RewriteInvocation) => unknown; + +/** A marker's rewrite behavior. Registered once by the consumer (for example + * ts-autocode-training registers `"use training"` with its capture interceptor); + * `"use "` in source is the shorthand that selects this configuration. */ +export interface RewriteConfig { + readonly marker: string; + readonly intercept?: RewriteInterceptor; +} type AnyMethod = (this: unknown, ...args: unknown[]) => unknown; -let interceptor: TrainableInterceptor | undefined; +const configs = new Map(); const swaps = new Map(); -/** One interceptor per process observes every trainable invocation - * (ts-autocode-training wires runtime capture here). */ -export function setTrainableInterceptor(next: TrainableInterceptor | undefined): void { - interceptor = next; +/** Normalizes a `"use "` directive to its canonical single-spaced form. */ +export function normalizeMarker(marker: string): string { + const trimmed = marker.trim().replace(/\s+/g, " "); + if (!/^use \S/.test(trimmed)) throw new TypeError(`rewrite marker must be a "use " directive: ${marker}`); + return trimmed; +} + +/** Single configuration entry point: binds a `"use "` marker to its rewrite + * behavior. After this, marking a method with that directive is all a consumer + * needs — weaving and swapping happen through the configured behavior, not + * through explicit `annotateRewrite`/`swapImplementation` calls. */ +export function configureRewrite(config: RewriteConfig): void { + const marker = normalizeMarker(config.marker); + configs.set(marker, Object.freeze({ ...config, marker })); +} + +export function rewriteMarkers(): readonly string[] { + return [...configs.keys()]; +} + +export function hasRewriteMarker(marker: string): boolean { + try { + return configs.has(normalizeMarker(marker)); + } catch { + return false; + } } -/** Hot-swappable advice: replaces the live implementation for a trainable id. - * Every woven call dispatches through the swap, without touching source. */ +/** Hot-swappable advice: replaces the live implementation for a rewrite id. + * Every woven call dispatches through the swap, without touching source. The + * single config entry point drives this on promotion; it is exported for tests + * and advanced orchestration, not the default consumer path. */ export function swapImplementation(id: string, implementation: AnyMethod): void { swaps.set(id, implementation); } @@ -49,9 +83,10 @@ export function swappedImplementation(id: string): AnyMethod | undefined { } /** Shared dispatch for the aspect and for wrapped free functions: hot-swap first, - * then the interceptor, then the original implementation. */ -export function dispatchTrainable( + * then the marker's configured interceptor, then the original implementation. */ +export function dispatchRewrite( id: string, + marker: string, methodName: string, original: AnyMethod, thisValue: unknown, @@ -61,17 +96,27 @@ export function dispatchTrainable( const active = swaps.get(id) ?? original; return active.apply(thisValue, next.length > 0 ? next : [...args]); }; - if (!interceptor) return proceed(); - return interceptor(Object.freeze({ id, methodName, thisValue, args, proceed })); + const config = configs.get(safeNormalize(marker)); + if (!config?.intercept) return proceed(); + return config.intercept(Object.freeze({ id, marker: safeNormalize(marker), methodName, thisValue, args, proceed })); +} + +function safeNormalize(marker: string): string { + try { + return normalizeMarker(marker); + } catch { + return marker; + } } -class TrainableAspectImpl { +class RewriteAspectImpl { intercept(context: AroundContext, joinpoint: JoinPoint, args: unknown[]): unknown { - const found = context.annotations(Trainable).find()[0]; + const found = context.annotations(Rewrite).find()[0]; const id = String(found?.args?.[0] ?? ""); + const marker = String(found?.args?.[1] ?? ""); const methodName = String((context.target as { propertyKey?: unknown }).propertyKey ?? id); const original: AnyMethod = (...next: unknown[]) => joinpoint(...next); - return dispatchTrainable(id, methodName, original, context.instance, args); + return dispatchRewrite(id, marker, methodName, original, context.instance, args); } } @@ -79,23 +124,26 @@ let wovenWeaver: unknown; /** Idempotent per weaver context; `configureTesting(WeaverModule)` swaps the * context, after which the next annotate re-enables the aspect. */ -export function enableTrainableWeaving(): void { +export function enableRewriteWeaving(): void { const weaver = getWeaver(); if (weaver === wovenWeaver) return; wovenWeaver = weaver; - applyLegacyDecorator(Around(on.methods.withAnnotations(Trainable)), TrainableAspectImpl.prototype, "intercept"); - const Enhanced = (Aspect()(TrainableAspectImpl) ?? TrainableAspectImpl) as typeof TrainableAspectImpl; + applyLegacyDecorator(Around(on.methods.withAnnotations(Rewrite)), RewriteAspectImpl.prototype, "intercept"); + const Enhanced = (Aspect()(RewriteAspectImpl) ?? RewriteAspectImpl) as typeof RewriteAspectImpl; weaver.enable(new Enhanced()); } const annotatedMethods = new WeakMap>(); -/** Weaves a class (or static) method for hot-swappable trainable dispatch. - * Walks the prototype chain to the owning container; idempotent per method. */ -export function annotateTrainable( +/** Weaves a class (or static) method for hot-swappable rewrite dispatch under a + * marker. Consumers do not call this directly: the `"use "` directive (via + * a consumer's discovery/register hook or decorator) is the shorthand that drives + * it. Walks the prototype chain to the owning container; idempotent per method. */ +export function annotateRewrite( owner: abstract new (...args: never[]) => unknown, methodName: string, id: string, + marker = "", ): void { const container = owningContainer(owner, methodName); if (!container) return; @@ -103,8 +151,16 @@ export function annotateTrainable( if (marked.has(methodName)) return; marked.add(methodName); annotatedMethods.set(container, marked); - enableTrainableWeaving(); - applyLegacyDecorator(Trainable(id) as LegacyMethodDecorator, container, methodName); + enableRewriteWeaving(); + applyLegacyDecorator(Rewrite(id, marker) as LegacyMethodDecorator, container, methodName); +} + +/** The prototype (or the constructor itself, for statics) that declares `methodName`. */ +export function declaringContainer( + owner: abstract new (...args: never[]) => unknown, + methodName: string, +): object | undefined { + return owningContainer(owner, methodName); } type LegacyMethodDecorator = ( diff --git a/packages/rewrite/src/index.ts b/packages/rewrite/src/index.ts index 58c078b..90ca801 100644 --- a/packages/rewrite/src/index.ts +++ b/packages/rewrite/src/index.ts @@ -7,13 +7,17 @@ export { promoteCandidate, revertPromotion } from "./promotion.js"; export type { PromotionResult, PromotionSnapshot, RewriteApproval } from "./promotion.js"; export { - Trainable, - annotateTrainable, - dispatchTrainable, - enableTrainableWeaving, + Rewrite, + annotateRewrite, + configureRewrite, + declaringContainer, + dispatchRewrite, + enableRewriteWeaving, + hasRewriteMarker, + normalizeMarker, restoreImplementation, - setTrainableInterceptor, + rewriteMarkers, swapImplementation, swappedImplementation, } from "./aspect.js"; -export type { TrainableInterceptor, TrainableInvocation } from "./aspect.js"; +export type { RewriteConfig, RewriteInterceptor, RewriteInvocation } from "./aspect.js"; diff --git a/packages/rewrite/test/weaving.test.ts b/packages/rewrite/test/weaving.test.ts index f690f00..bb637ee 100644 --- a/packages/rewrite/test/weaving.test.ts +++ b/packages/rewrite/test/weaving.test.ts @@ -3,21 +3,23 @@ import { WeaverModule } from "@aspectjs/core"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { - annotateTrainable, - dispatchTrainable, + annotateRewrite, + configureRewrite, + dispatchRewrite, restoreImplementation, - setTrainableInterceptor, swapImplementation, swappedImplementation, - type TrainableInvocation, + type RewriteInvocation, } from "../src/index.js"; +const MARKER = "use training"; + // Weaving is exercised entirely in memory: hot-swapped advice, never source // edits, so these tests can never rewrite themselves. -describe("trainable weaving", () => { +describe("rewrite weaving", () => { beforeEach(() => { configureTesting(WeaverModule); - setTrainableInterceptor(undefined); + configureRewrite({ marker: MARKER }); restoreImplementation("Router.route"); restoreImplementation("Router.fallback"); restoreImplementation("Static.echo"); @@ -29,12 +31,9 @@ describe("trainable weaving", () => { route(input: string): string { return input; } fallback(input: string): string { return input; } } - annotateTrainable(Router, "route", "Router.route"); + annotateRewrite(Router, "route", "Router.route", MARKER); const seen: string[] = []; - setTrainableInterceptor((invocation) => { - seen.push(invocation.id); - return invocation.proceed(); - }); + configureRewrite({ marker: MARKER, intercept: (invocation) => { seen.push(invocation.id); return invocation.proceed(); } }); const router = new Router(); expect(router.route("abc")).toBe("abc"); @@ -46,7 +45,7 @@ describe("trainable weaving", () => { class Router { route(input: string): string { return input; } } - annotateTrainable(Router, "route", "Router.route"); + annotateRewrite(Router, "route", "Router.route", MARKER); const router = new Router(); expect(router.route("abc")).toBe("abc"); @@ -63,29 +62,27 @@ describe("trainable weaving", () => { prefix = "id:"; route(input: string): string { return `${this.prefix}${input}`; } } - annotateTrainable(Router, "route", "Router.route"); - const invocations: TrainableInvocation[] = []; - setTrainableInterceptor((invocation) => { - invocations.push(invocation); - return invocation.proceed(); - }); + annotateRewrite(Router, "route", "Router.route", MARKER); + const invocations: RewriteInvocation[] = []; + configureRewrite({ marker: MARKER, intercept: (invocation) => { invocations.push(invocation); return invocation.proceed(); } }); swapImplementation("Router.route", function (this: unknown, input) { return `${(this as Router).prefix}${String(input).toUpperCase()}`; }); expect(new Router().route("abc")).toBe("id:ABC"); expect(invocations[0]?.id).toBe("Router.route"); + expect(invocations[0]?.marker).toBe(MARKER); expect(invocations[0]?.methodName).toBe("route"); expect(invocations[0]?.args).toEqual(["abc"]); }); - it("keeps swaps scoped to their trainable id", () => { + it("keeps swaps scoped to their rewrite id", () => { class Router { route(input: string): string { return input; } fallback(input: string): string { return input; } } - annotateTrainable(Router, "route", "Router.route"); - annotateTrainable(Router, "fallback", "Router.fallback"); + annotateRewrite(Router, "route", "Router.route", MARKER); + annotateRewrite(Router, "fallback", "Router.fallback", MARKER); swapImplementation("Router.route", () => "swapped"); const router = new Router(); @@ -97,20 +94,20 @@ describe("trainable weaving", () => { class Router { route(input: string): string { return input; } } - annotateTrainable(Router, "route", "Router.route"); - annotateTrainable(Router, "route", "Router.route"); - const interceptor = vi.fn((invocation: TrainableInvocation) => invocation.proceed()); - setTrainableInterceptor(interceptor); + annotateRewrite(Router, "route", "Router.route", MARKER); + annotateRewrite(Router, "route", "Router.route", MARKER); + const intercept = vi.fn((invocation: RewriteInvocation) => invocation.proceed()); + configureRewrite({ marker: MARKER, intercept }); expect(new Router().route("abc")).toBe("abc"); - expect(interceptor).toHaveBeenCalledTimes(1); + expect(intercept).toHaveBeenCalledTimes(1); }); it("weaves static methods and inherited methods through the owning container", () => { class Static { static echo(input: string): string { return input; } } - annotateTrainable(Static, "echo", "Static.echo"); + annotateRewrite(Static, "echo", "Static.echo", MARKER); swapImplementation("Static.echo", (input) => `static:${String(input)}`); expect(Static.echo("x")).toBe("static:x"); @@ -118,7 +115,7 @@ describe("trainable weaving", () => { route(input: string): string { return input; } } class Derived extends Base {} - annotateTrainable(Derived, "route", "Router.route"); + annotateRewrite(Derived, "route", "Router.route", MARKER); swapImplementation("Router.route", () => "woven"); expect(new Derived().route("abc")).toBe("woven"); expect(new Base().route("abc")).toBe("woven"); @@ -127,7 +124,7 @@ describe("trainable weaving", () => { it("dispatches wrapped free functions through the same swap registry", () => { const normalize = (input: string): string => input.trim(); const call = (input: string): unknown => - dispatchTrainable("free.normalize", "normalize", normalize as (...args: unknown[]) => unknown, undefined, [input]); + dispatchRewrite("free.normalize", MARKER, "normalize", normalize as (...args: unknown[]) => unknown, undefined, [input]); expect(call(" x ")).toBe("x"); swapImplementation("free.normalize", (input) => String(input).trim().toUpperCase()); @@ -136,19 +133,22 @@ describe("trainable weaving", () => { expect(call(" x ")).toBe("x"); }); - it("interceptors can observe without altering results, and clearing them restores plain dispatch", () => { + it("routes each marker to its own configured interceptor", () => { class Router { route(input: string): string { return input; } + ping(input: string): string { return input; } } - annotateTrainable(Router, "route", "Router.route"); - const router = new Router(); - const interceptor = vi.fn((invocation: TrainableInvocation) => invocation.proceed()); - setTrainableInterceptor(interceptor); - expect(router.route("abc")).toBe("abc"); - expect(interceptor).toHaveBeenCalledTimes(1); + const training: string[] = []; + const audit: string[] = []; + configureRewrite({ marker: "use training", intercept: (i) => { training.push(i.id); return i.proceed(); } }); + configureRewrite({ marker: "use audit", intercept: (i) => { audit.push(i.id); return i.proceed(); } }); + annotateRewrite(Router, "route", "Router.route", "use training"); + annotateRewrite(Router, "ping", "Router.ping", "use audit"); - setTrainableInterceptor(undefined); - expect(router.route("abc")).toBe("abc"); - expect(interceptor).toHaveBeenCalledTimes(1); + const router = new Router(); + router.route("a"); + router.ping("b"); + expect(training).toEqual(["Router.route"]); + expect(audit).toEqual(["Router.ping"]); }); }); diff --git a/packages/training/src/engine.ts b/packages/training/src/engine.ts index 1ac3feb..3e4fe90 100644 --- a/packages/training/src/engine.ts +++ b/packages/training/src/engine.ts @@ -50,12 +50,14 @@ export interface TrainingEngine { optimize(request: OptimizeRequest, context: EngineContext): Promise; } -/** Runs a proposed implementation against arguments in provider-owned isolation. */ +/** Runs a proposed implementation against arguments in provider-owned isolation. + * `receiver` is the live `this` when a hot-swapped instance method is invoked; + * sandboxed executors may ignore it. */ export type ImplementationExecutor = ( target: TrainableTarget, implementation: string, args: readonly unknown[], - options?: Readonly<{ timeoutMs?: number; signal?: AbortSignal }>, + options?: Readonly<{ timeoutMs?: number; signal?: AbortSignal; receiver?: unknown }>, ) => Promise; export async function optimizeCandidate( diff --git a/packages/training/src/training.ts b/packages/training/src/training.ts index 6a56b88..257ffbd 100644 --- a/packages/training/src/training.ts +++ b/packages/training/src/training.ts @@ -7,12 +7,13 @@ import { OpenInferenceSpanKind, SemanticConventions } from "@arizeai/openinferen import { SpanStatusCode, trace, type Attributes, type Span, type Tracer } from "@opentelemetry/api"; import { defineTrainingHarness, WriteAheadAgentBus, type JudgeRequest } from "ts-autocode-harness"; import { - annotateTrainable, - dispatchTrainable, + annotateRewrite, + configureRewrite, + declaringContainer, + dispatchRewrite, promoteCandidate, restoreImplementation, revertPromotion, - setTrainableInterceptor, swapImplementation, type PromotionResult, type PromotionSnapshot, @@ -45,6 +46,25 @@ import { const trainableAttribute = "ts_autocode.trainable.id"; +/** Training is one consumer of the generic rewrite engine; this is the marker it + * configures. The `"use training"` directive is the shorthand that weaves a method. */ +const trainingMarker = "use training"; + +// Register training's rewrite behavior once: every method woven under the +// "use training" marker routes through runtime capture, and `proceed` resolves +// the live (possibly hot-swapped) implementation so captures reflect what ran. +configureRewrite({ + marker: trainingMarker, + intercept: ({ id, methodName, thisValue, args, proceed }) => + runtime().invoke( + thisValue, + function (this: unknown, ...next: unknown[]) { return proceed(...next); }, + [...args], + defineTrainable(id), + methodName, + ), +}); + export interface CaptureSettings { readonly enabled?: boolean; readonly input?: boolean; @@ -462,8 +482,11 @@ class TrainingRuntime implements Training { if (!candidate.target.async) return; const executor = this.#settings.executor ?? defaultProviders.executor; if (!executor) return; - swapImplementation(candidate.trainableId, (...args: unknown[]) => - executor(candidate.target, candidate.implementation, args)); + // Normal function so the call receiver is captured and forwarded to + // executors that can bind it (the sandbox executor ignores it). + swapImplementation(candidate.trainableId, function (this: unknown, ...args: unknown[]) { + return executor(candidate.target, candidate.implementation, args, { receiver: this }); + }); } async flush(): Promise { @@ -631,22 +654,11 @@ export const training: Training = Object.freeze({ const wrappedMarker = Symbol.for("ts-autocode.wrapped"); -// Every woven or wrapped trainable dispatches through ts-autocode-rewrite's -// hot-swappable advice; this interceptor routes each invocation into runtime -// capture, and `proceed` resolves the live (possibly swapped) implementation. -setTrainableInterceptor(({ id, methodName, thisValue, args, proceed }) => - runtime().invoke( - thisValue, - function (this: unknown, ...next: unknown[]) { return proceed(...next); }, - [...args], - defineTrainable(id), - methodName, - )); - -/** Decorator form: `@trainable()`. Identity is inferred from the decorated class and - * method; pass a symbol (for example `defineTrainable("Router.route").symbol`) only - * to override the inferred id. The method is woven through the ts-autocode-rewrite - * aspect at first construction, so promoted candidates can hot-swap it. */ +/** Decorator form: `@trainable()`. Identity is inferred from the class that + * declares the method; pass a symbol (for example `defineTrainable("Router.route").symbol`) + * only to override the inferred id. The method is woven through the + * ts-autocode-rewrite aspect under the "use training" marker at first + * construction, so promoted candidates can hot-swap it. */ export function trainable(identity?: symbol): TrainableDecorator { if (identity !== undefined && typeof identity !== "symbol") { throw new TypeError("trainable identity must be a symbol; omit it to infer from the decorated method"); @@ -659,8 +671,10 @@ export function trainable(identity?: symbol): TrainableDecorator { const name = String(context.name); context.addInitializer(function (this: This) { const owner = (context.static ? this : (this as object).constructor) as abstract new (...args: never[]) => unknown; - const id = explicit?.id ?? `${inferredClassName(context.static ? owner : this) ?? "Anonymous"}.${name}`; - annotateTrainable(owner, name, id); + // Infer from the class that actually declares the method, so a base method + // first initialized through a subclass still resolves to Base.method. + const id = explicit?.id ?? `${declaringClassName(owner, name, context.static) ?? "Anonymous"}.${name}`; + annotateRewrite(owner, name, id, trainingMarker); }); return method; }; @@ -673,7 +687,7 @@ export function wrapTrainable unknown>(fn: F, id const name = fn.name || id; const method = fn as unknown as (this: unknown, ...args: unknown[]) => unknown; const wrapped = function (this: unknown, ...args: unknown[]): unknown { - return dispatchTrainable(id, name, method, this, args); + return dispatchRewrite(id, trainingMarker, name, method, this, args); }; Object.defineProperty(wrapped, "name", { value: name, configurable: true }); Object.defineProperty(wrapped, wrappedMarker, { value: true }); @@ -687,12 +701,18 @@ export function instrumentTrainable( methodName: string, id: string, ): void { - annotateTrainable(owner, methodName, id); + annotateRewrite(owner, methodName, id, trainingMarker); } -function inferredClassName(thisValue: unknown): string | undefined { - if (typeof thisValue === "function") return thisValue.name || undefined; - const constructor = (thisValue as { constructor?: unknown } | undefined)?.constructor; +/** Name of the class that declares `methodName`, walking to the owning prototype + * so an inherited method resolves to its base class rather than a subclass. */ +function declaringClassName( + owner: abstract new (...args: never[]) => unknown, + methodName: string, + isStatic: boolean, +): string | undefined { + const container = declaringContainer(owner, methodName); + const constructor = isStatic ? container : (container as { constructor?: unknown } | undefined)?.constructor; return typeof constructor === "function" && constructor.name ? constructor.name : undefined; }