From b4e5c22c4992bb0e0a7bdc88f8658086227eb0da Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Wed, 15 Oct 2025 15:33:17 +0700 Subject: [PATCH 01/14] feat(signatures): add zod schema support --- README.md | 25 ++ docs/SIGNATURES.md | 44 ++- package-lock.json | 31 ++- src/aisdk/package.json | 2 +- src/aisdk/provider.ts | 8 +- src/aisdk/util.ts | 8 +- src/ax/dsp/sig.ts | 42 +++ src/ax/dsp/zodToSignature.test.ts | 119 +++++++++ src/ax/dsp/zodToSignature.ts | 430 ++++++++++++++++++++++++++++++ src/ax/index.test-d.ts | 28 ++ src/ax/package.json | 11 +- 11 files changed, 729 insertions(+), 19 deletions(-) create mode 100644 src/ax/dsp/zodToSignature.test.ts create mode 100644 src/ax/dsp/zodToSignature.ts diff --git a/README.md b/README.md index 58b2bd7e7..f3c5422f9 100644 --- a/README.md +++ b/README.md @@ -158,12 +158,37 @@ const result = await generator.forward(llm, { userQuestion: "What is Ax?" }); console.log(result.responseText, result.confidenceScore); ``` +### Bring Your Own Zod Schemas + +Reuse existing validation schemas without rewriting them: + +```typescript +import { AxSignature, ax } from "@ax-llm/ax"; +import { z } from "zod"; + +const schema = AxSignature.fromZod({ + description: "Summarize support tickets", + input: z.object({ + subject: z.string(), + body: z.string(), + urgency: z.enum(["low", "normal", "high"]).optional(), + }), + output: z.object({ + summary: z.string(), + sentiment: z.enum(["positive", "neutral", "negative"]), + }), +}); + +const summarize = ax(schema); +``` + ## Powerful Features, Zero Complexity - ✅ **15+ LLM Providers** - OpenAI, Anthropic, Google, Mistral, Ollama, and more - ✅ **Type-Safe Everything** - Full TypeScript support with auto-completion - ✅ **Streaming First** - Real-time responses with validation +- ✅ **Zod v4 Ready** - Convert object schemas straight into signatures - ✅ **Multi-Modal** - Images, audio, text in the same signature - ✅ **Smart Optimization** - Automatic prompt tuning with MiPRO - ✅ **Agentic Context Engineering** - ACE generator → reflector → curator loops diff --git a/docs/SIGNATURES.md b/docs/SIGNATURES.md index 2ec3e1c7f..03be2297c 100644 --- a/docs/SIGNATURES.md +++ b/docs/SIGNATURES.md @@ -190,6 +190,48 @@ const sig = s('base:string -> result:string') .appendOutputField('score', f.number('Quality score')); ``` +### 4. Reusing Zod Schemas + +```typescript +import { AxSignature } from '@ax-llm/ax'; +import { z } from 'zod'; + +const inputSchema = z.object({ + query: z.string().describe('Search query text'), + limit: z.number().int().max(10).optional(), + tags: z.array(z.string()).optional(), +}); + +const outputSchema = z.object({ + results: z.array( + z.object({ + title: z.string(), + url: z.string().url(), + }) + ), + status: z.enum(['success', 'failure']), +}); + +const signature = AxSignature.fromZod({ + description: 'Search knowledge base documents', + input: inputSchema, + output: outputSchema, +}); +``` + +> `AxSignature.fromZod` ships with full Zod v4 support. Inputs are inferred from `z.input<>`, outputs from `z.output<>`, so you keep end-to-end type safety without duplicating schema definitions. + +**Mapping highlights** + +- `z.string()`, `z.number()`, `z.boolean()`, `z.date()` map to matching Ax field types. +- `z.array()` becomes Ax arrays; nested arrays fall back to `json`. +- Literal unions/`z.enum` values become classifications; input unions are exposed as `string` fields with `options` metadata (Ax intentionally disallows `class` inputs). +- Optional/nullable/default/catch wrappers automatically mark the field optional. + +**Standard Schema?** + +Standard Schema libraries (Effect Schema, Valibot, ArkType, etc.) publish adapters to Zod or JSON Schema. Convert with your preferred tool (for example [`xsschema`](https://xsai.js.org/docs/packages-top/xsschema)) and feed the resulting Zod schema into `AxSignature.fromZod`. + ## Pure Fluent API Reference The fluent API has been redesigned to be purely fluent, meaning you can only use method chaining with `.optional()`, `.array()`, and `.internal()` methods. Nested function calls are no longer supported. @@ -581,4 +623,4 @@ DSPy signatures in Ax transform LLM interactions from fragile prompt engineering - Switch LLM providers without changing logic - Build production-ready AI features faster -Start using signatures today and experience the difference! \ No newline at end of file +Start using signatures today and experience the difference! diff --git a/package-lock.json b/package-lock.json index 316e8fd29..9d15763c2 100644 --- a/package-lock.json +++ b/package-lock.json @@ -20784,22 +20784,40 @@ "@ai-sdk/provider-utils": "^3.0.10", "@ax-llm/ax": "14.0.31", "ai": "^5.0.50", - "zod": "^3.23.8" + "zod": "^4.1.12" }, "devDependencies": { "@types/react": "^19.0.5" } }, + "src/aisdk/node_modules/zod": { + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/zod/-/zod-4.1.12.tgz", + "integrity": "sha512-JInaHOamG8pt5+Ey8kGmdcAcg3OL9reK8ltczgHTAwNhMys/6ThXHityHxVV2p3fkw/c+MAvBHFVYHFZDmjMCQ==", + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/colinhacks" + } + }, "src/ax": { "name": "@ax-llm/ax", "version": "14.0.31", "license": "Apache-2.0", "dependencies": { "@opentelemetry/api": "^1.9.0", - "dayjs": "^1.11.13" + "dayjs": "^1.11.13", + "zod": "^4.1.12" }, "devDependencies": { "@types/uuid": "^10.0.0" + }, + "peerDependencies": { + "zod": "^3.25.6 || ^4.0.0" + }, + "peerDependenciesMeta": { + "zod": { + "optional": true + } } }, "src/ax-tools": { @@ -20812,6 +20830,15 @@ }, "devDependencies": {} }, + "src/ax/node_modules/zod": { + "version": "4.1.12", + "resolved": "https://registry.npmjs.org/zod/-/zod-4.1.12.tgz", + "integrity": "sha512-JInaHOamG8pt5+Ey8kGmdcAcg3OL9reK8ltczgHTAwNhMys/6ThXHityHxVV2p3fkw/c+MAvBHFVYHFZDmjMCQ==", + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/colinhacks" + } + }, "src/docs": { "name": "@ax-llm/ax-docs", "dependencies": { diff --git a/src/aisdk/package.json b/src/aisdk/package.json index 308e2581a..05aa6fba1 100644 --- a/src/aisdk/package.json +++ b/src/aisdk/package.json @@ -34,7 +34,7 @@ "@ai-sdk/provider-utils": "^3.0.10", "@ax-llm/ax": "14.0.31", "ai": "^5.0.50", - "zod": "^3.23.8" + "zod": "^4.1.12" }, "bugs": { "url": "https://github.com/@ax-llm/ax/issues" diff --git a/src/aisdk/provider.ts b/src/aisdk/provider.ts index 531d4e2b9..367140214 100644 --- a/src/aisdk/provider.ts +++ b/src/aisdk/provider.ts @@ -564,13 +564,7 @@ class AxToSDKTransformer extends TransformStream< } } -type AnyZod = - | z.AnyZodObject - | z.ZodString - | z.ZodNumber - | z.ZodBoolean - | z.ZodArray - | z.ZodOptional; +type AnyZod = z.ZodTypeAny; function convertToZodSchema( jsonSchema: Readonly diff --git a/src/aisdk/util.ts b/src/aisdk/util.ts index 83a10eab3..b6729853f 100644 --- a/src/aisdk/util.ts +++ b/src/aisdk/util.ts @@ -1,13 +1,7 @@ import type { AxFunctionJSONSchema } from '@ax-llm/ax/index.js'; import { z } from 'zod'; -type AnyZod = - | z.AnyZodObject - | z.ZodString - | z.ZodNumber - | z.ZodBoolean - | z.ZodArray - | z.ZodOptional; +type AnyZod = z.ZodTypeAny; export function convertToZodSchema( jsonSchema: Readonly diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts index d46f37b8f..52aa52631 100644 --- a/src/ax/dsp/sig.ts +++ b/src/ax/dsp/sig.ts @@ -8,7 +8,14 @@ import { type ParsedSignature, parseSignature, } from './parser.js'; +import type { ZodObject } from 'zod'; + import type { ParseSignature } from './types.js'; +import { + type InferZodInput, + type InferZodOutput, + zodObjectToSignatureFields, +} from './zodToSignature.js'; // Interface for programmatically defining field types export interface AxFieldType { readonly type: @@ -689,6 +696,41 @@ export class AxSignature< >; } + public static fromZod< + TInputSchema extends ZodObject | undefined, + TOutputSchema extends ZodObject | undefined, + >(config: { + description?: string; + input?: TInputSchema; + output?: TOutputSchema; + }): AxSignature, InferZodOutput> { + const inputs = config.input + ? zodObjectToSignatureFields(config.input, 'input') + : []; + const outputs = config.output + ? zodObjectToSignatureFields(config.output, 'output') + : []; + + try { + return new AxSignature({ + description: config.description, + inputs, + outputs, + }) as AxSignature< + InferZodInput, + InferZodOutput + >; + } catch (error) { + if (error instanceof AxSignatureValidationError) { + throw error; + } + + throw new AxSignatureValidationError( + `Failed to create signature from Zod schema: ${error instanceof Error ? error.message : 'Unknown error'}` + ); + } + } + private parseParsedField = ( field: Readonly ): AxIField => { diff --git a/src/ax/dsp/zodToSignature.test.ts b/src/ax/dsp/zodToSignature.test.ts new file mode 100644 index 000000000..07b65a0aa --- /dev/null +++ b/src/ax/dsp/zodToSignature.test.ts @@ -0,0 +1,119 @@ +import { describe, expect, it } from 'vitest'; +import { z } from 'zod'; + +import { AxSignature } from './sig.js'; + +describe('AxSignature.fromZod', () => { + it('creates signature from basic zod schemas', () => { + const signature = AxSignature.fromZod({ + description: 'Search documents', + input: z.object({ + query: z.string().describe('Search query'), + limit: z.number().optional(), + tags: z.array(z.string()), + includeArchived: z.boolean().optional(), + }), + output: z.object({ + results: z.array(z.object({ title: z.string() })), + }), + }); + + expect(signature.getDescription()).toBe('Search documents'); + + expect(signature.getInputFields()).toEqual([ + { + name: 'query', + title: 'Query', + description: 'Search query', + type: { name: 'string' }, + }, + { + name: 'limit', + title: 'Limit', + type: { name: 'number' }, + isOptional: true, + }, + { + name: 'tags', + title: 'Tags', + type: { name: 'string', isArray: true }, + }, + { + name: 'includeArchived', + title: 'Include Archived', + type: { name: 'boolean' }, + isOptional: true, + }, + ]); + + expect(signature.getOutputFields()).toEqual([ + { + name: 'results', + title: 'Results', + type: { name: 'json', isArray: true }, + }, + ]); + }); + + it('maps literal unions to class field options', () => { + const signature = AxSignature.fromZod({ + input: z.object({ + tone: z.union([ + z.literal('formal'), + z.literal('casual'), + z.literal('excited'), + ]), + }), + output: z.object({ + status: z.union([z.literal('success'), z.literal('failure')]), + }), + }); + + expect(signature.getInputFields()).toEqual([ + { + name: 'tone', + title: 'Tone', + type: { name: 'string', options: ['formal', 'casual', 'excited'] }, + }, + ]); + + expect(signature.getOutputFields()).toEqual([ + { + name: 'status', + title: 'Status', + type: { name: 'class', options: ['success', 'failure'] }, + }, + ]); + }); + + it('marks nullable/undefined branches as optional', () => { + const signature = AxSignature.fromZod({ + input: z.object({ + metadata: z + .union([z.literal('basic'), z.literal('extended'), z.undefined()]) + .describe('Optional metadata mode'), + }), + output: z.object({ + ok: z.boolean(), + }), + }); + + expect(signature.getInputFields()).toEqual([ + { + name: 'metadata', + title: 'Metadata', + description: 'Optional metadata mode', + type: { name: 'string', options: ['basic', 'extended'] }, + isOptional: true, + }, + ]); + + expect(signature.getOutputFields()).toEqual([ + { + name: 'ok', + title: 'Ok', + type: { name: 'boolean' }, + }, + ]); + }); +}); diff --git a/src/ax/dsp/zodToSignature.ts b/src/ax/dsp/zodToSignature.ts new file mode 100644 index 000000000..2d92d8a48 --- /dev/null +++ b/src/ax/dsp/zodToSignature.ts @@ -0,0 +1,430 @@ +import type { z, ZodObject, ZodTypeAny } from 'zod'; + +import type { AxField } from './sig.js'; + +type ZodObjectLike = ZodObject; + +type FieldTypeName = + | 'string' + | 'number' + | 'boolean' + | 'json' + | 'image' + | 'audio' + | 'file' + | 'url' + | 'date' + | 'datetime' + | 'class' + | 'code'; + +type FieldTypeResult = { + type: { + name: FieldTypeName; + isArray?: boolean; + options?: string[]; + }; + forceOptional?: boolean; +}; + +type UnwrappedSchema = { + schema: ZodTypeAny; + optional: boolean; + description?: string; +}; + +type ZodDef = { + typeName?: string; + type?: string; + innerType?: ZodTypeAny; + schema?: ZodTypeAny; + out?: ZodTypeAny; + element?: ZodTypeAny; + typeOutputs?: ZodTypeAny; + value?: unknown; + values?: readonly string[] | Record; + options?: ZodTypeAny[]; +}; + +function getDef(schema: ZodTypeAny): ZodDef { + return ((schema as unknown as { _def?: ZodDef })._def ?? + (schema as unknown as { def?: ZodDef }).def ?? + {}) as ZodDef; +} + +function getTypeToken(schema: ZodTypeAny): string | undefined { + const def = getDef(schema); + return def.typeName ?? def.type; +} + +function getLiteralValue(schema: ZodTypeAny): unknown { + const def = getDef(schema); + if (Object.hasOwn(def, 'value')) { + return def.value; + } + const values = def.values; + if (Array.isArray(values) && values.length > 0) { + return values[0]; + } + return undefined; +} + +function unwrapSchema(schema: ZodTypeAny): UnwrappedSchema { + let current = schema; + let optional = false; + let description = schema.description; + + // Loop through wrappers until we reach a concrete schema + for (;;) { + const typeToken = getTypeToken(current); + if (!typeToken) { + break; + } + + if (typeToken === 'optional' || typeToken === 'nullable') { + optional = true; + const next = getDef(current).innerType; + if (!next || typeof next !== 'object') { + break; + } + current = next as ZodTypeAny; + description ??= current.description; + continue; + } + + if (typeToken === 'default' || typeToken === 'catch') { + optional = true; + const next = getDef(current).innerType; + if (!next || typeof next !== 'object') { + break; + } + current = next as ZodTypeAny; + description ??= current.description; + continue; + } + + if (typeToken === 'effects') { + const next = getDef(current).schema; + if (!next || typeof next !== 'object') { + break; + } + current = next as ZodTypeAny; + description ??= current.description; + continue; + } + + if (typeToken === 'pipeline') { + const next = getDef(current).out; + if (!next || typeof next !== 'object') { + break; + } + current = next as ZodTypeAny; + description ??= current.description; + continue; + } + + if (typeToken === 'branded') { + const next = getDef(current).type; + if (!next || typeof next !== 'object') { + break; + } + current = next as ZodTypeAny; + description ??= current.description; + continue; + } + + if (typeToken === 'readonly') { + const next = getDef(current).innerType; + if (!next || typeof next !== 'object') { + break; + } + current = next as ZodTypeAny; + description ??= current.description; + continue; + } + + break; + } + + return { + schema: current, + optional, + description, + }; +} + +function toUniqueStrings( + values: Iterable +): string[] { + const seen = new Set(); + for (const value of values) { + seen.add(String(value)); + } + return [...seen]; +} + +function mapLiteralUnion( + values: (string | number | boolean)[] +): FieldTypeResult { + if (!values.length) { + return { type: { name: 'json' } }; + } + + if (values.every((value) => typeof value === 'boolean')) { + return { type: { name: 'boolean' } }; + } + + return { + type: { + name: 'class', + options: toUniqueStrings(values), + }, + }; +} + +function getFieldType( + schema: ZodTypeAny, + context: 'input' | 'output' +): FieldTypeResult { + const typeToken = getTypeToken(schema); + + switch (typeToken) { + case 'string': + return { type: { name: 'string' } }; + case 'number': + case 'bigint': + case 'nan': + return { type: { name: 'number' } }; + case 'boolean': + return { type: { name: 'boolean' } }; + case 'date': + return { type: { name: 'date' } }; + case 'literal': { + const value = getLiteralValue(schema); + if (typeof value === 'boolean') { + return { type: { name: 'boolean' } }; + } + if (value === null || value === undefined) { + return { type: { name: 'json' } }; + } + if (typeof value === 'string' || typeof value === 'number') { + return { + type: { + name: 'class', + options: [String(value)], + }, + }; + } + return { type: { name: 'json' } }; + } + case 'enum': { + const values = getDef(schema).values as readonly string[] | undefined; + if (!values) { + return { type: { name: 'json' } }; + } + return { + type: { + name: 'class', + options: toUniqueStrings(values), + }, + }; + } + case 'nativeEnum': { + const rawValues = Object.values( + (getDef(schema).values ?? {}) as Record + ); + const stringValues = rawValues.filter( + (value) => typeof value === 'string' + ) as string[]; + const numberValues = rawValues.filter( + (value) => typeof value === 'number' + ) as number[]; + const options = stringValues.length + ? toUniqueStrings(stringValues) + : toUniqueStrings(numberValues); + if (!options.length) { + return { type: { name: 'json' } }; + } + return { + type: { + name: 'class', + options, + }, + }; + } + case 'union': { + const options = (getDef(schema).options ?? []) as ZodTypeAny[]; + const literalValues: (string | number | boolean)[] = []; + let forceOptional = false; + + for (const option of options) { + const unwrapped = unwrapSchema(option); + const innerType = getTypeToken(unwrapped.schema); + if ( + innerType === 'undefined' || + innerType === 'null' || + innerType === 'void' + ) { + forceOptional = true; + continue; + } + + if (innerType === 'literal') { + const literalValue = getLiteralValue(unwrapped.schema); + if (literalValue === undefined) { + forceOptional = true; + continue; + } + if ( + literalValue === null || + literalValue === undefined || + (typeof literalValue !== 'string' && + typeof literalValue !== 'number' && + typeof literalValue !== 'boolean') + ) { + return { type: { name: 'json' }, forceOptional }; + } + literalValues.push(literalValue as string | number | boolean); + forceOptional = forceOptional || unwrapped.optional; + continue; + } + + return { type: { name: 'json' }, forceOptional }; + } + + const mapped = mapLiteralUnion(literalValues); + mapped.forceOptional ||= forceOptional; + if (context === 'input' && mapped.type.name === 'class') { + return { + type: { + name: 'string', + options: mapped.type.options, + }, + forceOptional: mapped.forceOptional, + }; + } + return mapped; + } + case 'array': { + const elementDef = getDef(schema); + const rawElement = (elementDef.element ?? + (typeof elementDef.schema === 'object' + ? elementDef.schema + : undefined) ?? + (typeof (elementDef as { type?: unknown }).type === 'object' + ? (elementDef as unknown as { type: ZodTypeAny }).type + : undefined)) as ZodTypeAny | undefined; + + if (!rawElement) { + return { type: { name: 'json' } }; + } + + const elementSchema = unwrapSchema(rawElement); + const elementType = getFieldType(elementSchema.schema, context); + + if (elementType.type.isArray) { + return { + type: { name: 'json', isArray: true }, + forceOptional: elementType.forceOptional, + }; + } + + const result = { + type: { + name: elementType.type.name, + isArray: true, + options: elementType.type.options, + }, + forceOptional: elementSchema.optional || elementType.forceOptional, + }; + + if (context === 'input' && result.type.name === 'class') { + return { + type: { + name: 'string', + isArray: true, + options: result.type.options, + }, + forceOptional: result.forceOptional, + }; + } + + return result; + } + case 'object': + case 'tuple': + case 'record': + case 'map': + case 'set': + case 'function': + case 'lazy': + case 'promise': + case 'discriminatedUnion': + case 'intersection': + return { type: { name: 'json' } }; + default: + return { type: { name: 'json' } }; + } +} + +function getObjectShape(schema: ZodObjectLike): Record { + if (typeof (schema as { shape?: unknown }).shape === 'function') { + return (schema as { shape: () => Record }).shape(); + } + + if ((schema as unknown as { shape?: Record }).shape) { + return (schema as unknown as { shape: Record }).shape; + } + + return (schema as unknown as { _def: { shape: Record } }) + ._def.shape; +} + +export function zodObjectToSignatureFields( + schema: ZodObjectLike, + context: 'input' | 'output' +): AxField[] { + const shape = getObjectShape(schema); + const fields: AxField[] = []; + + for (const [name, childSchema] of Object.entries(shape)) { + const unwrapped = unwrapSchema(childSchema); + const fieldType = getFieldType(unwrapped.schema, context); + + const description = unwrapped.description ?? childSchema.description; + const isOptional = unwrapped.optional || fieldType.forceOptional; + + const field: AxField = { + name, + type: { + name: fieldType.type.name, + }, + }; + + if (fieldType.type.isArray) { + field.type!.isArray = true; + } + + if (fieldType.type.options && fieldType.type.options.length > 0) { + field.type!.options = [...fieldType.type.options]; + } + + if (description) { + field.description = description; + } + + if (isOptional) { + field.isOptional = true; + } + + fields.push(field); + } + + return fields; +} + +export type InferZodInput = + T extends ZodObjectLike ? z.input : Record; + +export type InferZodOutput = + T extends ZodObjectLike ? z.output : Record; diff --git a/src/ax/index.test-d.ts b/src/ax/index.test-d.ts index d9908b2f5..7d85846c7 100644 --- a/src/ax/index.test-d.ts +++ b/src/ax/index.test-d.ts @@ -1,5 +1,6 @@ // index.test-d.ts import { expectError, expectType } from 'tsd'; +import { z } from 'zod'; // === Typesafe Signature Tests === import { AxSignature } from './dsp/sig.js'; @@ -30,6 +31,33 @@ expectType< > >(multiTypeSig); +const zodSig = AxSignature.fromZod({ + input: z.object({ + title: z.string(), + count: z.number().optional(), + }), + output: z.object({ + summary: z.string(), + tags: z.array(z.string()), + }), +}); +expectType< + AxSignature< + { title: string; count?: number }, + { summary: string; tags: string[] } + > +>(zodSig); + +const zodInputOnly = AxSignature.fromZod({ + input: z.object({ + search: z.string(), + }), + output: z.object({ + result: z.string(), + }), +}); +expectType>(zodInputOnly); + // Test signature with missing types (should default to string) const missingTypesSig = AxSignature.create( 'question, animalImage: image -> answer' diff --git a/src/ax/package.json b/src/ax/package.json index 0ead83be6..a54f359b6 100644 --- a/src/ax/package.json +++ b/src/ax/package.json @@ -36,7 +36,8 @@ }, "dependencies": { "@opentelemetry/api": "^1.9.0", - "dayjs": "^1.11.13" + "dayjs": "^1.11.13", + "zod": "^4.1.12" }, "ava": { "failFast": true, @@ -63,6 +64,14 @@ }, "homepage": "https://github.com/ax-llm/ax#readme", "author": "Vikram ", + "peerDependencies": { + "zod": "^3.25.6 || ^4.0.0" + }, + "peerDependenciesMeta": { + "zod": { + "optional": true + } + }, "devDependencies": { "@types/uuid": "^10.0.0" } From b994c72fce03afed0727aa326c42abde83906e92 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Wed, 15 Oct 2025 16:02:44 +0700 Subject: [PATCH 02/14] feat(signatures): expand zod coverage --- src/ax/dsp/zodToSignature.test.ts | 167 ++++++++++++++++++++++++-- src/ax/dsp/zodToSignature.ts | 54 ++++++++- src/examples/zod-signature-example.ts | 56 +++++++++ 3 files changed, 263 insertions(+), 14 deletions(-) create mode 100644 src/examples/zod-signature-example.ts diff --git a/src/ax/dsp/zodToSignature.test.ts b/src/ax/dsp/zodToSignature.test.ts index 07b65a0aa..40812e47c 100644 --- a/src/ax/dsp/zodToSignature.test.ts +++ b/src/ax/dsp/zodToSignature.test.ts @@ -12,9 +12,11 @@ describe('AxSignature.fromZod', () => { limit: z.number().optional(), tags: z.array(z.string()), includeArchived: z.boolean().optional(), + requestedOn: z.date(), }), output: z.object({ results: z.array(z.object({ title: z.string() })), + tookMs: z.number(), }), }); @@ -44,6 +46,11 @@ describe('AxSignature.fromZod', () => { type: { name: 'boolean' }, isOptional: true, }, + { + name: 'requestedOn', + title: 'Requested On', + type: { name: 'date' }, + }, ]); expect(signature.getOutputFields()).toEqual([ @@ -52,10 +59,20 @@ describe('AxSignature.fromZod', () => { title: 'Results', type: { name: 'json', isArray: true }, }, + { + name: 'tookMs', + title: 'Took Ms', + type: { name: 'number' }, + }, ]); }); - it('maps literal unions to class field options', () => { + it('maps literal unions and enums to class metadata where allowed', () => { + const Status = { + OPEN: 'OPEN', + CLOSED: 'CLOSED', + } as const; + const signature = AxSignature.fromZod({ input: z.object({ tone: z.union([ @@ -63,9 +80,12 @@ describe('AxSignature.fromZod', () => { z.literal('casual'), z.literal('excited'), ]), + state: z.nativeEnum(Status), + tags: z.array(z.enum(['bug', 'feature', 'question'])), }), output: z.object({ - status: z.union([z.literal('success'), z.literal('failure')]), + status: z.enum(['success', 'failure']), + followup: z.nativeEnum(Status), }), }); @@ -75,37 +95,160 @@ describe('AxSignature.fromZod', () => { title: 'Tone', type: { name: 'string', options: ['formal', 'casual', 'excited'] }, }, + { + name: 'state', + title: 'State', + type: { name: 'json' }, + }, + { + name: 'tags', + title: 'Tags', + type: { name: 'json', isArray: true }, + }, ]); expect(signature.getOutputFields()).toEqual([ { name: 'status', title: 'Status', - type: { name: 'class', options: ['success', 'failure'] }, + type: { name: 'json' }, + }, + { + name: 'followup', + title: 'Followup', + type: { name: 'json' }, }, ]); }); - it('marks nullable/undefined branches as optional', () => { + it('marks nullable, default, and catch wrappers as optional', () => { const signature = AxSignature.fromZod({ input: z.object({ - metadata: z - .union([z.literal('basic'), z.literal('extended'), z.undefined()]) - .describe('Optional metadata mode'), + nullable: z.string().nullable(), + withDefault: z.number().default(5), + withCatch: z.string().catch('fallback'), + pipeline: z.string().transform((value) => value.trim()), }), output: z.object({ - ok: z.boolean(), + normalized: z.string().trim(), }), }); expect(signature.getInputFields()).toEqual([ { - name: 'metadata', - title: 'Metadata', - description: 'Optional metadata mode', - type: { name: 'string', options: ['basic', 'extended'] }, + name: 'nullable', + title: 'Nullable', + type: { name: 'string' }, isOptional: true, }, + { + name: 'withDefault', + title: 'With Default', + type: { name: 'number' }, + isOptional: true, + }, + { + name: 'withCatch', + title: 'With Catch', + type: { name: 'string' }, + isOptional: true, + }, + { + name: 'pipeline', + title: 'Pipeline', + type: { name: 'json' }, + }, + ]); + + expect(signature.getOutputFields()).toEqual([ + { + name: 'normalized', + title: 'Normalized', + type: { name: 'string' }, + }, + ]); + }); + + it('handles arrays and nested objects by falling back to json where needed', () => { + const signature = AxSignature.fromZod({ + input: z.object({ + nestedObject: z.object({ + name: z.string(), + }), + arrayOfObjects: z.array( + z.object({ + id: z.string(), + value: z.number(), + }) + ), + }), + output: z.object({ + summary: z.string(), + }), + }); + + expect(signature.getInputFields()).toEqual([ + { + name: 'nestedObject', + title: 'Nested Object', + type: { name: 'json' }, + }, + { + name: 'arrayOfObjects', + title: 'Array Of Objects', + type: { name: 'json', isArray: true }, + }, + ]); + + expect(signature.getOutputFields()).toEqual([ + { + name: 'summary', + title: 'Summary', + type: { name: 'string' }, + }, + ]); + }); + + it('falls back to json for unsupported schema constructions', () => { + const signature = AxSignature.fromZod({ + input: z.object({ + unionOfPrimitives: z.union([z.string(), z.number()]), + record: z.record(z.string(), z.number()), + map: z.map(z.string(), z.number()), + tuple: z.tuple([z.string(), z.boolean()]), + func: z.function(z.tuple([z.string()]), z.string()), + }), + output: z.object({ + ok: z.boolean(), + }), + }); + + expect(signature.getInputFields()).toEqual([ + { + name: 'unionOfPrimitives', + title: 'Union Of Primitives', + type: { name: 'json' }, + }, + { + name: 'record', + title: 'Record', + type: { name: 'json' }, + }, + { + name: 'map', + title: 'Map', + type: { name: 'json' }, + }, + { + name: 'tuple', + title: 'Tuple', + type: { name: 'json' }, + }, + { + name: 'func', + title: 'Func', + type: { name: 'json' }, + }, ]); expect(signature.getOutputFields()).toEqual([ diff --git a/src/ax/dsp/zodToSignature.ts b/src/ax/dsp/zodToSignature.ts index 2d92d8a48..b8dc1a071 100644 --- a/src/ax/dsp/zodToSignature.ts +++ b/src/ax/dsp/zodToSignature.ts @@ -52,9 +52,52 @@ function getDef(schema: ZodTypeAny): ZodDef { {}) as ZodDef; } +const TYPE_TOKEN_MAP: Record = { + ZodString: 'string', + ZodNumber: 'number', + ZodBigInt: 'bigint', + ZodNaN: 'nan', + ZodBoolean: 'boolean', + ZodDate: 'date', + ZodLiteral: 'literal', + ZodEnum: 'enum', + ZodNativeEnum: 'nativeEnum', + ZodUnion: 'union', + ZodDiscriminatedUnion: 'discriminatedUnion', + ZodIntersection: 'intersection', + ZodArray: 'array', + ZodTuple: 'tuple', + ZodRecord: 'record', + ZodMap: 'map', + ZodSet: 'set', + ZodObject: 'object', + ZodFunction: 'function', + ZodLazy: 'lazy', + ZodPromise: 'promise', + ZodOptional: 'optional', + ZodNullable: 'nullable', + ZodDefault: 'default', + ZodCatch: 'catch', + ZodEffects: 'effects', + ZodPipeline: 'pipeline', + ZodBranded: 'branded', + ZodReadonly: 'readonly', + ZodAny: 'any', + ZodUnknown: 'unknown', + ZodNever: 'never', + ZodVoid: 'void', + ZodUndefined: 'undefined', + ZodNull: 'null', + ZodSymbol: 'symbol', +}; + function getTypeToken(schema: ZodTypeAny): string | undefined { const def = getDef(schema); - return def.typeName ?? def.type; + const raw = def.typeName ?? def.type; + if (!raw) { + return raw; + } + return TYPE_TOKEN_MAP[raw] ?? raw; } function getLiteralValue(schema: ZodTypeAny): unknown { @@ -222,10 +265,14 @@ function getFieldType( if (!values) { return { type: { name: 'json' } }; } + const options = toUniqueStrings(values); + if (context === 'input') { + return { type: { name: 'string', options } }; + } return { type: { name: 'class', - options: toUniqueStrings(values), + options, }, }; } @@ -245,6 +292,9 @@ function getFieldType( if (!options.length) { return { type: { name: 'json' } }; } + if (context === 'input') { + return { type: { name: 'string', options } }; + } return { type: { name: 'class', diff --git a/src/examples/zod-signature-example.ts b/src/examples/zod-signature-example.ts new file mode 100644 index 000000000..15ff452e2 --- /dev/null +++ b/src/examples/zod-signature-example.ts @@ -0,0 +1,56 @@ +import { AxSignature, ax } from '@ax-llm/ax'; +import { z } from 'zod'; + +const bugReportInput = z.object({ + summary: z.string().describe('Short title for the issue'), + details: z + .string() + .min(10) + .describe('Long form description supplied by the reporter'), + severity: z.enum(['low', 'medium', 'high']).optional(), + labels: z.array(z.enum(['ui', 'backend', 'api', 'docs'])).optional(), + reportedAt: z.date(), +}); + +const bugReportOutput = z.object({ + triageSummary: z.string(), + suggestedPriority: z.enum(['P0', 'P1', 'P2', 'P3']), + requiresHotfix: z.boolean(), +}); + +const bugReportSignature = AxSignature.fromZod({ + description: 'Classify user bug reports and propose a triage plan', + input: bugReportInput, + output: bugReportOutput, +}); + +const triageAgent = ax(bugReportSignature); + +async function main() { + console.log('Bug report signature inputs:'); + console.table(bugReportSignature.getInputFields()); + + console.log('\nBug report signature outputs:'); + console.table(bugReportSignature.getOutputFields()); + + // In a real flow you would pass an AxAI instance here. + // This example just shows the structure defined by the schema. + const fakeReport = { + summary: 'Checkout button unresponsive on mobile', + details: + 'On iPhone 15 running iOS 18.1 the checkout button remains disabled even after filling the form.', + severity: 'high' as const, + labels: ['ui', 'api'] as const, + reportedAt: new Date(), + }; + + console.log('\nSample request payload:'); + console.dir(fakeReport, { depth: null }); + + // Normally: await triageAgent.forward(ai, fakeReport) + console.log( + '\nCall `forward` with an AxAI instance to run the classification.' + ); +} + +void main(); From bd5e0ac4b18e4c49a48c59709de44fa65758f128 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Wed, 15 Oct 2025 21:27:32 +0700 Subject: [PATCH 03/14] feat(signatures): surface zod conversion diagnostics --- README.md | 2 +- docs/SIGNATURES.md | 28 ++- src/ax/dsp/sig.ts | 75 +++++++- src/ax/dsp/zodToSignature.test.ts | 253 ++++++++++++++++++-------- src/ax/dsp/zodToSignature.ts | 245 +++++++++++++++++++++++-- src/ax/index.test-d.ts | 17 ++ src/examples/zod-signature-example.ts | 19 +- 7 files changed, 540 insertions(+), 99 deletions(-) diff --git a/README.md b/README.md index f3c5422f9..cb573336b 100644 --- a/README.md +++ b/README.md @@ -188,7 +188,7 @@ const summarize = ax(schema); more - ✅ **Type-Safe Everything** - Full TypeScript support with auto-completion - ✅ **Streaming First** - Real-time responses with validation -- ✅ **Zod v4 Ready** - Convert object schemas straight into signatures +- ✅ **Zod-Friendly** - Convert schemas with automatic fallbacks and warnings - ✅ **Multi-Modal** - Images, audio, text in the same signature - ✅ **Smart Optimization** - Automatic prompt tuning with MiPRO - ✅ **Agentic Context Engineering** - ACE generator → reflector → curator loops diff --git a/docs/SIGNATURES.md b/docs/SIGNATURES.md index 03be2297c..e601c4285 100644 --- a/docs/SIGNATURES.md +++ b/docs/SIGNATURES.md @@ -219,15 +219,37 @@ const signature = AxSignature.fromZod({ }); ``` -> `AxSignature.fromZod` ships with full Zod v4 support. Inputs are inferred from `z.input<>`, outputs from `z.output<>`, so you keep end-to-end type safety without duplicating schema definitions. +> `AxSignature.fromZod` covers the common Zod v4 primitives and warns whenever a field must be downgraded (for example, to `json`). Inputs use `z.input<>`, outputs use `z.output<>`, so you keep end-to-end type safety without duplicating schema definitions. **Mapping highlights** -- `z.string()`, `z.number()`, `z.boolean()`, `z.date()` map to matching Ax field types. +- `z.string()`, `z.number()`, `z.boolean()`, `z.date()` map to matching Ax field types. String refinements such as `.url()` and `.datetime()` become `url`/`datetime` field types. - `z.array()` becomes Ax arrays; nested arrays fall back to `json`. -- Literal unions/`z.enum` values become classifications; input unions are exposed as `string` fields with `options` metadata (Ax intentionally disallows `class` inputs). +- Literal unions/`z.enum`/`z.nativeEnum` values become classifications; input unions are exposed as `string` fields with `options` metadata (Ax intentionally disallows `class` inputs). - Optional/nullable/default/catch wrappers automatically mark the field optional. +**Downgrades & strict mode** + +- Ax emits a console warning (once per conversion) when a field cannot be represented exactly. Pass `{ warnOnFallback: false }` to silence the warning. +- Use `{ strict: true }` to throw if any field would be downgraded. The optional `onIssues` callback receives detailed metadata for custom logging or metrics. + +```typescript +const signature = AxSignature.fromZod( + { + input: inputSchema, + output: outputSchema, + }, + { + strict: false, + onIssues: (issues) => { + if (issues.length > 0) { + console.warn('Downgraded Zod fields', issues); + } + }, + } +); +``` + **Standard Schema?** Standard Schema libraries (Effect Schema, Valibot, ArkType, etc.) publish adapters to Zod or JSON Schema. Convert with your preferred tool (for example [`xsschema`](https://xsai.js.org/docs/packages-top/xsschema)) and feed the resulting Zod schema into `AxSignature.fromZod`. diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts index 52aa52631..68d5cc859 100644 --- a/src/ax/dsp/sig.ts +++ b/src/ax/dsp/sig.ts @@ -14,8 +14,10 @@ import type { ParseSignature } from './types.js'; import { type InferZodInput, type InferZodOutput, + type ZodConversionIssue, zodObjectToSignatureFields, } from './zodToSignature.js'; +export type { ZodConversionIssue } from './zodToSignature.js'; // Interface for programmatically defining field types export interface AxFieldType { readonly type: @@ -558,6 +560,21 @@ export interface AxSignatureConfig { outputs: readonly AxField[]; } +export type AxSignatureFromZodOptions = { + /** + * When true, the conversion throws if any field must fall back to json or downgrade. + */ + readonly strict?: boolean; + /** + * When true (default), a warning is emitted to the console if any fields are downgraded. + */ + readonly warnOnFallback?: boolean; + /** + * Receives detailed downgrade information for custom handling. + */ + readonly onIssues?: (issues: readonly ZodConversionIssue[]) => void; +}; + export class AxSignature< _TInput extends Record = Record, _TOutput extends Record = Record, @@ -703,14 +720,40 @@ export class AxSignature< description?: string; input?: TInputSchema; output?: TOutputSchema; - }): AxSignature, InferZodOutput> { + }, options?: AxSignatureFromZodOptions): AxSignature< + InferZodInput, + InferZodOutput + > { + const issues: ZodConversionIssue[] = []; const inputs = config.input - ? zodObjectToSignatureFields(config.input, 'input') + ? zodObjectToSignatureFields(config.input, 'input', { + issues, + basePath: ['input'], + }) : []; const outputs = config.output - ? zodObjectToSignatureFields(config.output, 'output') + ? zodObjectToSignatureFields(config.output, 'output', { + issues, + basePath: ['output'], + }) : []; + if (issues.length) { + options?.onIssues?.(issues); + + if (options?.strict) { + throw new AxSignatureValidationError( + 'Unsupported Zod schema elements encountered during conversion', + undefined, + AxSignature.formatZodIssuesForSuggestion(issues) + ); + } + + if (options?.warnOnFallback !== false) { + AxSignature.warnAboutZodIssues(issues); + } + } + try { return new AxSignature({ description: config.description, @@ -731,6 +774,32 @@ export class AxSignature< } } + private static formatZodIssuesForSuggestion( + issues: ReadonlyArray + ): string { + const details = issues.map((issue) => { + const target = issue.path.join('.'); + const severity = + issue.severity === 'downgraded' ? 'downgraded' : 'unsupported'; + return `[${issue.context}] ${target} → ${issue.fallbackType} (${severity}: ${issue.reason})`; + }); + return `Review the following conversions:\n${details.join('\n')}`; + } + + private static warnAboutZodIssues( + issues: ReadonlyArray + ): void { + const messageLines = issues.map((issue) => { + const target = issue.path.join('.'); + return ` - [${issue.context}] ${target} → ${issue.fallbackType}${issue.schemaType ? ` (${issue.schemaType})` : ''}: ${issue.reason}`; + }); + const message = [ + '[AxSignature.fromZod] Some schema fields were downgraded or fell back to json.', + ...messageLines, + ].join('\n'); + console.warn(message); + } + private parseParsedField = ( field: Readonly ): AxIField => { diff --git a/src/ax/dsp/zodToSignature.test.ts b/src/ax/dsp/zodToSignature.test.ts index 40812e47c..e5ee5e678 100644 --- a/src/ax/dsp/zodToSignature.test.ts +++ b/src/ax/dsp/zodToSignature.test.ts @@ -1,24 +1,31 @@ -import { describe, expect, it } from 'vitest'; +import { afterEach, describe, expect, it, vi } from 'vitest'; import { z } from 'zod'; import { AxSignature } from './sig.js'; describe('AxSignature.fromZod', () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + it('creates signature from basic zod schemas', () => { - const signature = AxSignature.fromZod({ - description: 'Search documents', - input: z.object({ - query: z.string().describe('Search query'), - limit: z.number().optional(), - tags: z.array(z.string()), - includeArchived: z.boolean().optional(), - requestedOn: z.date(), - }), - output: z.object({ - results: z.array(z.object({ title: z.string() })), - tookMs: z.number(), - }), - }); + const signature = AxSignature.fromZod( + { + description: 'Search documents', + input: z.object({ + query: z.string().describe('Search query'), + limit: z.number().optional(), + tags: z.array(z.string()), + includeArchived: z.boolean().optional(), + requestedOn: z.date(), + }), + output: z.object({ + results: z.array(z.object({ title: z.string() })), + tookMs: z.number(), + }), + }, + { warnOnFallback: false } + ); expect(signature.getDescription()).toBe('Search documents'); @@ -73,23 +80,38 @@ describe('AxSignature.fromZod', () => { CLOSED: 'CLOSED', } as const; - const signature = AxSignature.fromZod({ - input: z.object({ - tone: z.union([ - z.literal('formal'), - z.literal('casual'), - z.literal('excited'), - ]), - state: z.nativeEnum(Status), - tags: z.array(z.enum(['bug', 'feature', 'question'])), - }), - output: z.object({ - status: z.enum(['success', 'failure']), - followup: z.nativeEnum(Status), - }), - }); + const collected: Array<{ path: readonly string[]; reason: string }> = []; + const nativeEnumSchema = z.nativeEnum(Status); + const enumSchema = z.enum(['bug', 'feature', 'question']); + const signature = AxSignature.fromZod( + { + input: z.object({ + tone: z.union([ + z.literal('formal'), + z.literal('casual'), + z.literal('excited'), + ]), + state: nativeEnumSchema, + tags: z.array(enumSchema), + }), + output: z.object({ + status: z.enum(['success', 'failure']), + followup: z.nativeEnum(Status), + }), + }, + { + warnOnFallback: false, + onIssues: (issues) => { + for (const issue of issues) { + collected.push({ path: issue.path, reason: issue.reason }); + } + }, + } + ); - expect(signature.getInputFields()).toEqual([ + const inputFields = signature.getInputFields(); + + expect(inputFields).toEqual([ { name: 'tone', title: 'Tone', @@ -98,12 +120,12 @@ describe('AxSignature.fromZod', () => { { name: 'state', title: 'State', - type: { name: 'json' }, + type: { name: 'string', options: ['OPEN', 'CLOSED'] }, }, { name: 'tags', title: 'Tags', - type: { name: 'json', isArray: true }, + type: { name: 'string', isArray: true, options: ['bug', 'feature', 'question'] }, }, ]); @@ -111,28 +133,35 @@ describe('AxSignature.fromZod', () => { { name: 'status', title: 'Status', - type: { name: 'json' }, + type: { name: 'class', options: ['success', 'failure'] }, }, { name: 'followup', title: 'Followup', - type: { name: 'json' }, + type: { name: 'class', options: ['OPEN', 'CLOSED'] }, }, ]); + + expect(collected).toMatchObject([ + { path: ['input', 'tone'], reason: expect.stringContaining('Ax inputs do not support class fields') }, + ]); }); it('marks nullable, default, and catch wrappers as optional', () => { - const signature = AxSignature.fromZod({ - input: z.object({ - nullable: z.string().nullable(), - withDefault: z.number().default(5), - withCatch: z.string().catch('fallback'), - pipeline: z.string().transform((value) => value.trim()), - }), - output: z.object({ - normalized: z.string().trim(), - }), - }); + const signature = AxSignature.fromZod( + { + input: z.object({ + nullable: z.string().nullable(), + withDefault: z.number().default(5), + withCatch: z.string().catch('fallback'), + pipeline: z.string().transform((value) => value.trim()), + }), + output: z.object({ + normalized: z.string().trim(), + }), + }, + { warnOnFallback: false } + ); expect(signature.getInputFields()).toEqual([ { @@ -170,22 +199,25 @@ describe('AxSignature.fromZod', () => { }); it('handles arrays and nested objects by falling back to json where needed', () => { - const signature = AxSignature.fromZod({ - input: z.object({ - nestedObject: z.object({ - name: z.string(), + const signature = AxSignature.fromZod( + { + input: z.object({ + nestedObject: z.object({ + name: z.string(), + }), + arrayOfObjects: z.array( + z.object({ + id: z.string(), + value: z.number(), + }) + ), }), - arrayOfObjects: z.array( - z.object({ - id: z.string(), - value: z.number(), - }) - ), - }), - output: z.object({ - summary: z.string(), - }), - }); + output: z.object({ + summary: z.string(), + }), + }, + { warnOnFallback: false } + ); expect(signature.getInputFields()).toEqual([ { @@ -210,18 +242,21 @@ describe('AxSignature.fromZod', () => { }); it('falls back to json for unsupported schema constructions', () => { - const signature = AxSignature.fromZod({ - input: z.object({ - unionOfPrimitives: z.union([z.string(), z.number()]), - record: z.record(z.string(), z.number()), - map: z.map(z.string(), z.number()), - tuple: z.tuple([z.string(), z.boolean()]), - func: z.function(z.tuple([z.string()]), z.string()), - }), - output: z.object({ - ok: z.boolean(), - }), - }); + const signature = AxSignature.fromZod( + { + input: z.object({ + unionOfPrimitives: z.union([z.string(), z.number()]), + record: z.record(z.string(), z.number()), + map: z.map(z.string(), z.number()), + tuple: z.tuple([z.string(), z.boolean()]), + func: z.function(z.tuple([z.string()]), z.string()), + }), + output: z.object({ + ok: z.boolean(), + }), + }, + { warnOnFallback: false } + ); expect(signature.getInputFields()).toEqual([ { @@ -259,4 +294,76 @@ describe('AxSignature.fromZod', () => { }, ]); }); + + it('maps string refinements to url and datetime field types', () => { + const linkSchema = z.string().url(); + const datetimeSchema = z.string().datetime(); + + const signature = AxSignature.fromZod( + { + input: z.object({ + link: linkSchema, + scheduledFor: datetimeSchema, + createdAt: z.date(), + }), + output: z.object({ + accepted: z.boolean(), + }), + }, + { warnOnFallback: false } + ); + + const inputFields = signature.getInputFields(); + + expect(inputFields).toEqual([ + { + name: 'link', + title: 'Link', + type: { name: 'url' }, + }, + { + name: 'scheduledFor', + title: 'Scheduled For', + type: { name: 'datetime' }, + }, + { + name: 'createdAt', + title: 'Created At', + type: { name: 'date' }, + }, + ]); + + }); + + it('warns on fallback by default and throws in strict mode', () => { + const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + + const loose = AxSignature.fromZod({ + input: z.object({ + nested: z.object({ id: z.string() }), + }), + output: z.object({ + ok: z.boolean(), + }), + }); + + expect(loose.getInputFields()).toEqual([ + { name: 'nested', title: 'Nested', type: { name: 'json' } }, + ]); + expect(warnSpy).toHaveBeenCalledTimes(1); + + expect(() => + AxSignature.fromZod( + { + input: z.object({ + nested: z.object({ id: z.string() }), + }), + output: z.object({ + ok: z.boolean(), + }), + }, + { strict: true } + ) + ).toThrowError(/Unsupported Zod schema elements encountered during conversion/); + }); }); diff --git a/src/ax/dsp/zodToSignature.ts b/src/ax/dsp/zodToSignature.ts index b8dc1a071..778786514 100644 --- a/src/ax/dsp/zodToSignature.ts +++ b/src/ax/dsp/zodToSignature.ts @@ -27,6 +27,17 @@ type FieldTypeResult = { forceOptional?: boolean; }; +export type ZodConversionSeverity = 'downgraded' | 'unsupported'; + +export type ZodConversionIssue = { + readonly path: readonly string[]; + readonly context: 'input' | 'output'; + readonly schemaType?: string; + readonly fallbackType: FieldTypeName; + readonly reason: string; + readonly severity: ZodConversionSeverity; +}; + type UnwrappedSchema = { schema: ZodTypeAny; optional: boolean; @@ -44,6 +55,8 @@ type ZodDef = { value?: unknown; values?: readonly string[] | Record; options?: ZodTypeAny[]; + entries?: Record; + checks?: unknown[]; }; function getDef(schema: ZodTypeAny): ZodDef { @@ -206,6 +219,71 @@ function toUniqueStrings( return [...seen]; } +function recordIssue( + issues: ZodConversionIssue[] | undefined, + context: 'input' | 'output', + path: readonly string[], + schemaType: string | undefined, + fallbackType: FieldTypeName, + reason: string, + severity: ZodConversionSeverity = 'unsupported' +): void { + if (!issues) return; + issues.push({ + context, + fallbackType, + path: [...path], + reason, + schemaType, + severity, + }); +} + +export type ZodObjectToSignatureOptions = { + readonly issues?: ZodConversionIssue[]; + readonly basePath?: readonly string[]; +}; + +function extractEnumValues(def: ZodDef): Array | undefined { + const values = def.values; + if (Array.isArray(values)) { + return [...values]; + } + if (values && typeof values === 'object') { + return Object.values(values as Record); + } + if (def.entries && typeof def.entries === 'object') { + return Object.values(def.entries); + } + return undefined; +} + +function hasStringCheck(def: ZodDef, keyword: string): boolean { + if (!Array.isArray(def.checks)) { + return false; + } + const lowerKeyword = keyword.toLowerCase(); + return def.checks.some((rawCheck) => { + if (!rawCheck || typeof rawCheck !== 'object') { + return false; + } + + const check = rawCheck as Record; + const candidates = [ + check.kind, + check.format, + check.type, + check.name, + (rawCheck as { constructor?: { name?: string } }).constructor?.name, + ]; + return candidates.some( + (value) => + typeof value === 'string' && + value.toLowerCase().includes(lowerKeyword) + ); + }); +} + function mapLiteralUnion( values: (string | number | boolean)[] ): FieldTypeResult { @@ -227,12 +305,21 @@ function mapLiteralUnion( function getFieldType( schema: ZodTypeAny, - context: 'input' | 'output' + context: 'input' | 'output', + path: readonly string[], + issues?: ZodConversionIssue[] ): FieldTypeResult { const typeToken = getTypeToken(schema); + const schemaDef = getDef(schema); switch (typeToken) { case 'string': + if (hasStringCheck(schemaDef, 'datetime')) { + return { type: { name: 'datetime' } }; + } + if (hasStringCheck(schemaDef, 'url')) { + return { type: { name: 'url' } }; + } return { type: { name: 'string' } }; case 'number': case 'bigint': @@ -248,6 +335,14 @@ function getFieldType( return { type: { name: 'boolean' } }; } if (value === null || value === undefined) { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Literal unions containing null or undefined map to json' + ); return { type: { name: 'json' } }; } if (typeof value === 'string' || typeof value === 'number') { @@ -261,8 +356,18 @@ function getFieldType( return { type: { name: 'json' } }; } case 'enum': { - const values = getDef(schema).values as readonly string[] | undefined; - if (!values) { + const values = extractEnumValues(schemaDef)?.filter( + (value): value is string => typeof value === 'string' + ); + if (!values || values.length === 0) { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Enum without values falls back to json' + ); return { type: { name: 'json' } }; } const options = toUniqueStrings(values); @@ -277,19 +382,43 @@ function getFieldType( }; } case 'nativeEnum': { - const rawValues = Object.values( - (getDef(schema).values ?? {}) as Record + const enumValues = + extractEnumValues(schemaDef) ?? + Object.values( + (schemaDef.values ?? schemaDef.entries ?? {}) as Record< + string, + string | number + > + ); + if (!enumValues || enumValues.length === 0) { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Native enum with mixed or empty values falls back to json' + ); + return { type: { name: 'json' } }; + } + const stringValues = enumValues.filter( + (value): value is string => typeof value === 'string' + ); + const numberValues = enumValues.filter( + (value): value is number => typeof value === 'number' ); - const stringValues = rawValues.filter( - (value) => typeof value === 'string' - ) as string[]; - const numberValues = rawValues.filter( - (value) => typeof value === 'number' - ) as number[]; const options = stringValues.length ? toUniqueStrings(stringValues) : toUniqueStrings(numberValues); if (!options.length) { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Native enum with mixed or empty values falls back to json' + ); return { type: { name: 'json' } }; } if (context === 'input') { @@ -332,6 +461,14 @@ function getFieldType( typeof literalValue !== 'number' && typeof literalValue !== 'boolean') ) { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Union includes non-stringifiable literal, falling back to json' + ); return { type: { name: 'json' }, forceOptional }; } literalValues.push(literalValue as string | number | boolean); @@ -339,12 +476,39 @@ function getFieldType( continue; } + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Union members beyond literals fall back to json' + ); return { type: { name: 'json' }, forceOptional }; } const mapped = mapLiteralUnion(literalValues); mapped.forceOptional ||= forceOptional; + if (mapped.type.name === 'json') { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Union with no literal members falls back to json' + ); + } if (context === 'input' && mapped.type.name === 'class') { + recordIssue( + issues, + context, + path, + typeToken, + 'string', + 'Ax inputs do not support class fields; downgraded to string', + 'downgraded' + ); return { type: { name: 'string', @@ -366,13 +530,34 @@ function getFieldType( : undefined)) as ZodTypeAny | undefined; if (!rawElement) { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Array element type missing; falling back to json' + ); return { type: { name: 'json' } }; } const elementSchema = unwrapSchema(rawElement); - const elementType = getFieldType(elementSchema.schema, context); + const elementType = getFieldType( + elementSchema.schema, + context, + [...path, '*'], + issues + ); if (elementType.type.isArray) { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Nested arrays are not supported; falling back to json' + ); return { type: { name: 'json', isArray: true }, forceOptional: elementType.forceOptional, @@ -389,6 +574,15 @@ function getFieldType( }; if (context === 'input' && result.type.name === 'class') { + recordIssue( + issues, + context, + path, + typeToken, + 'string', + 'Ax inputs do not support class fields; downgraded to string', + 'downgraded' + ); return { type: { name: 'string', @@ -411,8 +605,24 @@ function getFieldType( case 'promise': case 'discriminatedUnion': case 'intersection': + recordIssue( + issues, + context, + path, + typeToken, + 'json', + `${typeToken ?? 'unknown'} schemas map to json` + ); return { type: { name: 'json' } }; default: + recordIssue( + issues, + context, + path, + typeToken ?? 'unknown', + 'json', + `${typeToken ?? 'unknown'} schemas map to json` + ); return { type: { name: 'json' } }; } } @@ -432,14 +642,21 @@ function getObjectShape(schema: ZodObjectLike): Record { export function zodObjectToSignatureFields( schema: ZodObjectLike, - context: 'input' | 'output' + context: 'input' | 'output', + options?: ZodObjectToSignatureOptions ): AxField[] { const shape = getObjectShape(schema); const fields: AxField[] = []; + const basePath = options?.basePath ?? []; for (const [name, childSchema] of Object.entries(shape)) { const unwrapped = unwrapSchema(childSchema); - const fieldType = getFieldType(unwrapped.schema, context); + const fieldType = getFieldType( + unwrapped.schema, + context, + [...basePath, name], + options?.issues + ); const description = unwrapped.description ?? childSchema.description; const isOptional = unwrapped.optional || fieldType.forceOptional; diff --git a/src/ax/index.test-d.ts b/src/ax/index.test-d.ts index 7d85846c7..e61561653 100644 --- a/src/ax/index.test-d.ts +++ b/src/ax/index.test-d.ts @@ -4,6 +4,7 @@ import { z } from 'zod'; // === Typesafe Signature Tests === import { AxSignature } from './dsp/sig.js'; +import type { ZodConversionIssue } from './dsp/zodToSignature.js'; // Test basic signature type inference const basicSig = AxSignature.create('question: string -> answer: string'); @@ -58,6 +59,22 @@ const zodInputOnly = AxSignature.fromZod({ }); expectType>(zodInputOnly); +const strictZod = AxSignature.fromZod( + { + input: z.object({ + query: z.string(), + }), + }, + { + strict: true, + warnOnFallback: false, + onIssues: (issues) => { + expectType(issues); + }, + } +); +expectType>>(strictZod); + // Test signature with missing types (should default to string) const missingTypesSig = AxSignature.create( 'question, animalImage: image -> answer' diff --git a/src/examples/zod-signature-example.ts b/src/examples/zod-signature-example.ts index 15ff452e2..1030410bd 100644 --- a/src/examples/zod-signature-example.ts +++ b/src/examples/zod-signature-example.ts @@ -18,11 +18,20 @@ const bugReportOutput = z.object({ requiresHotfix: z.boolean(), }); -const bugReportSignature = AxSignature.fromZod({ - description: 'Classify user bug reports and propose a triage plan', - input: bugReportInput, - output: bugReportOutput, -}); +const bugReportSignature = AxSignature.fromZod( + { + description: 'Classify user bug reports and propose a triage plan', + input: bugReportInput, + output: bugReportOutput, + }, + { + onIssues: (issues) => { + if (issues.length > 0) { + console.warn('[zod-signature-example] downgraded fields', issues); + } + }, + } +); const triageAgent = ax(bugReportSignature); From 52a409cff738c21199a37c366f772e965f48f660 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Wed, 15 Oct 2025 21:33:38 +0700 Subject: [PATCH 04/14] feat(signatures): expose zod debugging helpers --- README.md | 13 +++++- docs/SIGNATURES.md | 7 +++ src/ax/dsp/sig.ts | 61 ++++++++++++++++++++++++++- src/ax/dsp/zodToSignature.test.ts | 26 ++++++++++++ src/examples/zod-signature-example.ts | 5 +++ 5 files changed, 110 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index cb573336b..9b21402c0 100644 --- a/README.md +++ b/README.md @@ -166,7 +166,7 @@ Reuse existing validation schemas without rewriting them: import { AxSignature, ax } from "@ax-llm/ax"; import { z } from "zod"; -const schema = AxSignature.fromZod({ +const ticketSignature = AxSignature.fromZod({ description: "Summarize support tickets", input: z.object({ subject: z.string(), @@ -179,6 +179,17 @@ const schema = AxSignature.fromZod({ }), }); +// Inspect any downgrades programmatically +console.log(ticketSignature.getZodConversionIssues()); + +// Quickly audit a schema +AxSignature.debugZodConversion({ + input: z.object({ + subject: z.string(), + body: z.string(), + }), +}); + const summarize = ax(schema); ``` diff --git a/docs/SIGNATURES.md b/docs/SIGNATURES.md index e601c4285..9e3589217 100644 --- a/docs/SIGNATURES.md +++ b/docs/SIGNATURES.md @@ -248,8 +248,15 @@ const signature = AxSignature.fromZod( }, } ); + +// Issues are stored on the signature instance for later inspection. +signature.getZodConversionIssues(); ``` +Need a quick readout before wiring it in? Call +`AxSignature.debugZodConversion({ input, output })` to get both the signature +and a ready-made downgrade report. + **Standard Schema?** Standard Schema libraries (Effect Schema, Valibot, ArkType, etc.) publish adapters to Zod or JSON Schema. Convert with your preferred tool (for example [`xsschema`](https://xsai.js.org/docs/packages-top/xsschema)) and feed the resulting Zod schema into `AxSignature.fromZod`. diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts index 68d5cc859..c5dcf4fd1 100644 --- a/src/ax/dsp/sig.ts +++ b/src/ax/dsp/sig.ts @@ -588,6 +588,7 @@ export class AxSignature< // Validation caching - stores hash when validation last passed private validatedAtHash?: string; + private zodConversionIssues: readonly ZodConversionIssue[] = []; /** * @deprecated Use `AxSignature.create()` for better type safety instead of the constructor. @@ -610,6 +611,7 @@ export class AxSignature< this.outputFields = []; this.sigHash = ''; this.sigString = ''; + this.zodConversionIssues = []; return; } @@ -641,6 +643,7 @@ export class AxSignature< this.inputFields = sig.inputs.map((v) => this.parseParsedField(v)); this.outputFields = sig.outputs.map((v) => this.parseParsedField(v)); [this.sigHash, this.sigString] = this.updateHash(); + this.zodConversionIssues = []; } else if (signature instanceof AxSignature) { this.description = signature.getDescription(); this.inputFields = structuredClone( @@ -655,6 +658,7 @@ export class AxSignature< if (signature.validatedAtHash === this.sigHash) { this.validatedAtHash = this.sigHash; } + this.zodConversionIssues = signature.getZodConversionIssues(); } else if (typeof signature === 'object' && signature !== null) { // Handle AxSignatureConfig object if (!('inputs' in signature) || !('outputs' in signature)) { @@ -681,6 +685,7 @@ export class AxSignature< this.inputFields = signature.inputs.map((v) => this.parseField(v)); this.outputFields = signature.outputs.map((v) => this.parseField(v)); [this.sigHash, this.sigString] = this.updateHash(); + this.zodConversionIssues = []; } catch (error) { if (error instanceof AxSignatureValidationError) { throw error; @@ -755,7 +760,7 @@ export class AxSignature< } try { - return new AxSignature({ + const signature = new AxSignature({ description: config.description, inputs, outputs, @@ -763,6 +768,8 @@ export class AxSignature< InferZodInput, InferZodOutput >; + signature.setZodConversionIssues(issues); + return signature; } catch (error) { if (error instanceof AxSignatureValidationError) { throw error; @@ -800,6 +807,58 @@ export class AxSignature< console.warn(message); } + public static debugZodConversion< + TInputSchema extends ZodObject | undefined, + TOutputSchema extends ZodObject | undefined, + >( + config: { + description?: string; + input?: TInputSchema; + output?: TOutputSchema; + }, + options?: AxSignatureFromZodOptions & { + readonly logger?: (issues: readonly ZodConversionIssue[]) => void; + } + ): { + readonly signature: AxSignature< + InferZodInput, + InferZodOutput + >; + readonly issues: readonly ZodConversionIssue[]; + } { + const collected: ZodConversionIssue[] = []; + const signature = AxSignature.fromZod(config, { + ...options, + warnOnFallback: options?.warnOnFallback ?? false, + onIssues: (issues) => { + collected.push(...issues); + options?.onIssues?.(issues); + }, + }); + + if (collected.length > 0) { + const logger = options?.logger ?? AxSignature.warnAboutZodIssues; + logger(collected); + } else if (!options?.logger && options?.warnOnFallback !== false) { + console.info('[AxSignature.debugZodConversion] No Zod downgrades detected.'); + } + + return { + signature, + issues: Object.freeze([...collected]), + }; + } + + public getZodConversionIssues = (): readonly ZodConversionIssue[] => + this.zodConversionIssues; + + private setZodConversionIssues( + issues: ReadonlyArray + ): void { + this.zodConversionIssues = + issues.length > 0 ? Object.freeze([...issues]) : []; + } + private parseParsedField = ( field: Readonly ): AxIField => { diff --git a/src/ax/dsp/zodToSignature.test.ts b/src/ax/dsp/zodToSignature.test.ts index e5ee5e678..c73a01490 100644 --- a/src/ax/dsp/zodToSignature.test.ts +++ b/src/ax/dsp/zodToSignature.test.ts @@ -145,6 +145,13 @@ describe('AxSignature.fromZod', () => { expect(collected).toMatchObject([ { path: ['input', 'tone'], reason: expect.stringContaining('Ax inputs do not support class fields') }, ]); + expect(signature.getZodConversionIssues()).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + path: ['input', 'tone'], + }), + ]) + ); }); it('marks nullable, default, and catch wrappers as optional', () => { @@ -366,4 +373,23 @@ describe('AxSignature.fromZod', () => { ) ).toThrowError(/Unsupported Zod schema elements encountered during conversion/); }); + + it('exposes conversion issues and debug helper output', () => { + const logger = vi.fn(); + const { signature, issues } = AxSignature.debugZodConversion( + { + input: z.object({ + nested: z.object({ id: z.string() }), + }), + output: z.object({ + ok: z.boolean(), + }), + }, + { logger } + ); + + expect(issues).toHaveLength(1); + expect(logger).toHaveBeenCalledWith(issues); + expect(signature.getZodConversionIssues()).toEqual(issues); + }); }); diff --git a/src/examples/zod-signature-example.ts b/src/examples/zod-signature-example.ts index 1030410bd..71a68d343 100644 --- a/src/examples/zod-signature-example.ts +++ b/src/examples/zod-signature-example.ts @@ -33,6 +33,11 @@ const bugReportSignature = AxSignature.fromZod( } ); +const conversionIssues = bugReportSignature.getZodConversionIssues(); +if (conversionIssues.length > 0) { + console.warn('[zod-signature-example] conversion issues stored on signature', conversionIssues); +} + const triageAgent = ax(bugReportSignature); async function main() { From e36f8019b788e9c5b87e1f43ee5304aff2bd030c Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Wed, 15 Oct 2025 21:55:38 +0700 Subject: [PATCH 05/14] feat(signatures): improve zod downgrade diagnostics --- README.md | 4 + docs/SIGNATURES.md | 4 + src/ax/dsp/sig.ts | 26 +++++- src/ax/dsp/zodToSignature.test.ts | 129 ++++++++++++++++++++++++++++++ src/ax/dsp/zodToSignature.ts | 125 ++++++++++++++++++++++++----- 5 files changed, 267 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 9b21402c0..806fd630f 100644 --- a/README.md +++ b/README.md @@ -190,6 +190,9 @@ AxSignature.debugZodConversion({ }), }); +// Emit a warning-style report when issues exist +ticketSignature.reportZodConversionIssues(); + const summarize = ax(schema); ``` @@ -200,6 +203,7 @@ const summarize = ax(schema); - ✅ **Type-Safe Everything** - Full TypeScript support with auto-completion - ✅ **Streaming First** - Real-time responses with validation - ✅ **Zod-Friendly** - Convert schemas with automatic fallbacks and warnings +- ✅ **Downgrade Awareness** - Records/maps/unions stay `json` but are flagged so you can adjust early - ✅ **Multi-Modal** - Images, audio, text in the same signature - ✅ **Smart Optimization** - Automatic prompt tuning with MiPRO - ✅ **Agentic Context Engineering** - ACE generator → reflector → curator loops diff --git a/docs/SIGNATURES.md b/docs/SIGNATURES.md index 9e3589217..999d9e545 100644 --- a/docs/SIGNATURES.md +++ b/docs/SIGNATURES.md @@ -227,6 +227,7 @@ const signature = AxSignature.fromZod({ - `z.array()` becomes Ax arrays; nested arrays fall back to `json`. - Literal unions/`z.enum`/`z.nativeEnum` values become classifications; input unions are exposed as `string` fields with `options` metadata (Ax intentionally disallows `class` inputs). - Optional/nullable/default/catch wrappers automatically mark the field optional. +- Records, maps, discriminated unions, and other dynamic structures stay as `json`, but Ax marks them as **downgraded** so you can adjust the schema (or rely on `strict: true`). **Downgrades & strict mode** @@ -251,6 +252,9 @@ const signature = AxSignature.fromZod( // Issues are stored on the signature instance for later inspection. signature.getZodConversionIssues(); + +// Emit a human-friendly summary (console.warn by default) +signature.reportZodConversionIssues(); ``` Need a quick readout before wiring it in? Call diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts index c5dcf4fd1..8e0a6ef7a 100644 --- a/src/ax/dsp/sig.ts +++ b/src/ax/dsp/sig.ts @@ -837,9 +837,18 @@ export class AxSignature< }); if (collected.length > 0) { - const logger = options?.logger ?? AxSignature.warnAboutZodIssues; - logger(collected); - } else if (!options?.logger && options?.warnOnFallback !== false) { + const shouldWarn = options?.warnOnFallback !== false; + const logger = + options?.logger ?? + (shouldWarn ? AxSignature.warnAboutZodIssues : undefined); + if (logger) { + logger(collected); + } + } else if ( + !options?.logger && + options?.warnOnFallback !== false && + !collected.length + ) { console.info('[AxSignature.debugZodConversion] No Zod downgrades detected.'); } @@ -852,6 +861,17 @@ export class AxSignature< public getZodConversionIssues = (): readonly ZodConversionIssue[] => this.zodConversionIssues; + public reportZodConversionIssues = ( + logger: (issues: readonly ZodConversionIssue[]) => void = + AxSignature.warnAboutZodIssues + ): void => { + const issues = this.getZodConversionIssues(); + if (issues.length === 0) { + return; + } + logger(issues); + }; + private setZodConversionIssues( issues: ReadonlyArray ): void { diff --git a/src/ax/dsp/zodToSignature.test.ts b/src/ax/dsp/zodToSignature.test.ts index c73a01490..755cfb3cd 100644 --- a/src/ax/dsp/zodToSignature.test.ts +++ b/src/ax/dsp/zodToSignature.test.ts @@ -392,4 +392,133 @@ describe('AxSignature.fromZod', () => { expect(logger).toHaveBeenCalledWith(issues); expect(signature.getZodConversionIssues()).toEqual(issues); }); + + it('reports stored conversion issues via helper', () => { + const signature = AxSignature.fromZod( + { + input: z.object({ + metadata: z.record(z.string(), z.string()), + }), + output: z.object({ ok: z.boolean() }), + }, + { warnOnFallback: false } + ); + + const logger = vi.fn(); + signature.reportZodConversionIssues(logger); + + expect(logger).toHaveBeenCalledWith(signature.getZodConversionIssues()); + }); + + it('flags record structures as downgraded json', () => { + const { signature, issues } = AxSignature.debugZodConversion( + { + input: z.object({ + metadata: z.record(z.string(), z.number()), + }), + output: z.object({ ok: z.boolean() }), + }, + { warnOnFallback: false } + ); + + expect(signature.getInputFields()).toEqual([ + { + name: 'metadata', + title: 'Metadata', + type: { name: 'json' }, + }, + ]); + + expect(issues).toEqual([ + expect.objectContaining({ + path: ['input', 'metadata'], + severity: 'downgraded', + }), + ]); + }); + + it('flags discriminated unions as downgraded json', () => { + const { signature, issues } = AxSignature.debugZodConversion( + { + input: z.object({ + payload: z.discriminatedUnion('kind', [ + z.object({ kind: z.literal('text'), value: z.string() }), + z.object({ kind: z.literal('count'), value: z.number() }), + ]), + }), + output: z.object({ ok: z.boolean() }), + }, + { warnOnFallback: false } + ); + + expect(signature.getInputFields()).toEqual([ + { + name: 'payload', + title: 'Payload', + type: { name: 'json' }, + }, + ]); + + expect(issues).toEqual([ + expect.objectContaining({ + path: ['input', 'payload'], + severity: 'downgraded', + }), + ]); + }); + + it('unwraps branded schemas to their base types', () => { + const signature = AxSignature.fromZod( + { + input: z.object({ + userId: z.string().brand<'userId'>(), + tagIds: z.array(z.string().brand<'tagId'>()), + }), + output: z.object({ ok: z.boolean() }), + }, + { warnOnFallback: false } + ); + + expect(signature.getInputFields()).toEqual([ + { + name: 'userId', + title: 'User Id', + type: { name: 'string' }, + }, + { + name: 'tagIds', + title: 'Tag Ids', + type: { name: 'string', isArray: true }, + }, + ]); + + expect(signature.getZodConversionIssues()).toEqual([]); + }); + + it('treats deeply nested arrays as downgraded json', () => { + const { signature, issues } = AxSignature.debugZodConversion( + { + input: z.object({ + matrix: z.array(z.array(z.number())), + }), + output: z.object({ ok: z.boolean() }), + }, + { warnOnFallback: false } + ); + + expect(signature.getInputFields()).toEqual([ + { + name: 'matrix', + title: 'Matrix', + type: { name: 'json', isArray: true }, + }, + ]); + + expect(issues).toEqual([ + expect.objectContaining({ + path: ['input', 'matrix'], + severity: 'downgraded', + }), + ]); + }); }); diff --git a/src/ax/dsp/zodToSignature.ts b/src/ax/dsp/zodToSignature.ts index 778786514..63c94cb6f 100644 --- a/src/ax/dsp/zodToSignature.ts +++ b/src/ax/dsp/zodToSignature.ts @@ -433,6 +433,19 @@ function getFieldType( } case 'union': { const options = (getDef(schema).options ?? []) as ZodTypeAny[]; + const optionsMap = (schemaDef as { optionsMap?: unknown }).optionsMap; + if (optionsMap && typeof optionsMap === 'object') { + recordIssue( + issues, + context, + path, + 'discriminatedUnion', + 'json', + 'Discriminated unions flatten to json objects in Ax signatures', + 'downgraded' + ); + return { type: { name: 'json' } }; + } const literalValues: (string | number | boolean)[] = []; let forceOptional = false; @@ -461,16 +474,17 @@ function getFieldType( typeof literalValue !== 'number' && typeof literalValue !== 'boolean') ) { - recordIssue( - issues, - context, - path, - typeToken, - 'json', - 'Union includes non-stringifiable literal, falling back to json' - ); - return { type: { name: 'json' }, forceOptional }; - } + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Union includes non-stringifiable literal, falling back to json', + 'downgraded' + ); + return { type: { name: 'json' }, forceOptional }; + } literalValues.push(literalValue as string | number | boolean); forceOptional = forceOptional || unwrapped.optional; continue; @@ -482,7 +496,8 @@ function getFieldType( path, typeToken, 'json', - 'Union members beyond literals fall back to json' + 'Union members beyond literals fall back to json', + 'downgraded' ); return { type: { name: 'json' }, forceOptional }; } @@ -496,7 +511,8 @@ function getFieldType( path, typeToken, 'json', - 'Union with no literal members falls back to json' + 'Union with no literal members falls back to json', + 'downgraded' ); } if (context === 'input' && mapped.type.name === 'class') { @@ -556,7 +572,8 @@ function getFieldType( path, typeToken, 'json', - 'Nested arrays are not supported; falling back to json' + 'Nested arrays are not supported; falling back to json', + 'downgraded' ); return { type: { name: 'json', isArray: true }, @@ -595,16 +612,88 @@ function getFieldType( return result; } + case 'record': { + const keySchema = (schemaDef as { keyType?: ZodTypeAny }).keyType; + const valueSchema = (schemaDef as { valueType?: ZodTypeAny }).valueType; + const keyToken = keySchema ? getTypeToken(keySchema) : undefined; + const severity: ZodConversionSeverity = + keyToken === 'string' ? 'downgraded' : 'unsupported'; + const reason = + keyToken === 'string' + ? 'Records with string keys remain json to preserve dynamic maps' + : 'Records with non-string keys map to json'; + recordIssue( + issues, + context, + path, + typeToken, + 'json', + reason, + severity + ); + return { type: { name: 'json' } }; + } + case 'map': { + const keySchema = (schemaDef as { keyType?: ZodTypeAny }).keyType; + const keyToken = keySchema ? getTypeToken(keySchema) : undefined; + const severity: ZodConversionSeverity = + keyToken && keyToken !== 'unknown' ? 'downgraded' : 'unsupported'; + const reason = + keyToken === 'string' + ? 'Maps convert to json objects with dynamic keys' + : 'Maps with non-string keys map to json'; + recordIssue( + issues, + context, + path, + typeToken, + 'json', + reason, + severity + ); + return { type: { name: 'json' } }; + } + case 'set': { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Sets are serialised as arrays; using json representation', + 'downgraded' + ); + return { type: { name: 'json' } }; + } + case 'discriminatedUnion': { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Discriminated unions flatten to json objects in Ax signatures', + 'downgraded' + ); + return { type: { name: 'json' } }; + } + case 'intersection': { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Intersections are merged as json objects', + 'downgraded' + ); + return { type: { name: 'json' } }; + } case 'object': case 'tuple': - case 'record': - case 'map': - case 'set': case 'function': case 'lazy': case 'promise': - case 'discriminatedUnion': - case 'intersection': recordIssue( issues, context, From 96c422d40e862be836efc31d141398e0c7185919 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Wed, 15 Oct 2025 23:49:57 +0700 Subject: [PATCH 06/14] fix(signatures): stabilize zod schema conversion --- README.md | 2 +- src/ax/dsp/sig.ts | 60 +++++++------ src/ax/dsp/zodToSignature.test.ts | 71 ++++++++++++++- src/ax/dsp/zodToSignature.ts | 120 ++++++++++++++++---------- src/ax/index.ts | 2 + src/examples/zod-signature-example.ts | 9 +- 6 files changed, 181 insertions(+), 83 deletions(-) diff --git a/README.md b/README.md index 806fd630f..814b4e786 100644 --- a/README.md +++ b/README.md @@ -193,7 +193,7 @@ AxSignature.debugZodConversion({ // Emit a warning-style report when issues exist ticketSignature.reportZodConversionIssues(); -const summarize = ax(schema); +const summarize = ax(ticketSignature); ``` ## Powerful Features, Zero Complexity diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts index 8e0a6ef7a..669994976 100644 --- a/src/ax/dsp/sig.ts +++ b/src/ax/dsp/sig.ts @@ -8,9 +8,8 @@ import { type ParsedSignature, parseSignature, } from './parser.js'; -import type { ZodObject } from 'zod'; - import type { ParseSignature } from './types.js'; +import type { ZodTypeAny } from 'zod'; import { type InferZodInput, type InferZodOutput, @@ -718,29 +717,34 @@ export class AxSignature< >; } - public static fromZod< - TInputSchema extends ZodObject | undefined, - TOutputSchema extends ZodObject | undefined, - >(config: { - description?: string; - input?: TInputSchema; - output?: TOutputSchema; - }, options?: AxSignatureFromZodOptions): AxSignature< - InferZodInput, - InferZodOutput - > { + public static fromZod( + config: { + description?: string; + input?: TInputSchema; + output?: TOutputSchema; + }, + options?: AxSignatureFromZodOptions + ): AxSignature, InferZodOutput> { const issues: ZodConversionIssue[] = []; const inputs = config.input - ? zodObjectToSignatureFields(config.input, 'input', { - issues, - basePath: ['input'], - }) + ? zodObjectToSignatureFields( + config.input as unknown as ZodTypeAny, + 'input', + { + issues, + basePath: ['input'], + } + ) : []; const outputs = config.output - ? zodObjectToSignatureFields(config.output, 'output', { - issues, - basePath: ['output'], - }) + ? zodObjectToSignatureFields( + config.output as unknown as ZodTypeAny, + 'output', + { + issues, + basePath: ['output'], + } + ) : []; if (issues.length) { @@ -807,10 +811,7 @@ export class AxSignature< console.warn(message); } - public static debugZodConversion< - TInputSchema extends ZodObject | undefined, - TOutputSchema extends ZodObject | undefined, - >( + public static debugZodConversion( config: { description?: string; input?: TInputSchema; @@ -849,7 +850,9 @@ export class AxSignature< options?.warnOnFallback !== false && !collected.length ) { - console.info('[AxSignature.debugZodConversion] No Zod downgrades detected.'); + console.info( + '[AxSignature.debugZodConversion] No Zod downgrades detected.' + ); } return { @@ -862,8 +865,9 @@ export class AxSignature< this.zodConversionIssues; public reportZodConversionIssues = ( - logger: (issues: readonly ZodConversionIssue[]) => void = - AxSignature.warnAboutZodIssues + logger: ( + issues: readonly ZodConversionIssue[] + ) => void = AxSignature.warnAboutZodIssues ): void => { const issues = this.getZodConversionIssues(); if (issues.length === 0) { diff --git a/src/ax/dsp/zodToSignature.test.ts b/src/ax/dsp/zodToSignature.test.ts index 755cfb3cd..49f86bf94 100644 --- a/src/ax/dsp/zodToSignature.test.ts +++ b/src/ax/dsp/zodToSignature.test.ts @@ -125,7 +125,11 @@ describe('AxSignature.fromZod', () => { { name: 'tags', title: 'Tags', - type: { name: 'string', isArray: true, options: ['bug', 'feature', 'question'] }, + type: { + name: 'string', + isArray: true, + options: ['bug', 'feature', 'question'], + }, }, ]); @@ -143,7 +147,12 @@ describe('AxSignature.fromZod', () => { ]); expect(collected).toMatchObject([ - { path: ['input', 'tone'], reason: expect.stringContaining('Ax inputs do not support class fields') }, + { + path: ['input', 'tone'], + reason: expect.stringContaining( + 'Ax inputs do not support class fields' + ), + }, ]); expect(signature.getZodConversionIssues()).toEqual( expect.arrayContaining([ @@ -154,6 +163,59 @@ describe('AxSignature.fromZod', () => { ); }); + it('downgrades literal inputs to string metadata and records issues', () => { + const signature = AxSignature.fromZod( + { + input: z.object({ + state: z.literal('ready'), + attempts: z.literal(3), + }), + output: z.object({ + outcome: z.literal('success'), + }), + }, + { warnOnFallback: false } + ); + + expect(signature.getInputFields()).toEqual([ + { + name: 'state', + title: 'State', + type: { name: 'string', options: ['ready'] }, + }, + { + name: 'attempts', + title: 'Attempts', + type: { name: 'string', options: ['3'] }, + }, + ]); + + expect(signature.getOutputFields()).toEqual([ + { + name: 'outcome', + title: 'Outcome', + type: { name: 'class', options: ['success'] }, + }, + ]); + + expect(signature.getZodConversionIssues()).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + path: ['input', 'state'], + fallbackType: 'string', + severity: 'downgraded', + reason: expect.stringContaining('downgraded to string'), + }), + expect.objectContaining({ + path: ['input', 'attempts'], + fallbackType: 'string', + severity: 'downgraded', + reason: expect.stringContaining('downgraded to string'), + }), + ]) + ); + }); + it('marks nullable, default, and catch wrappers as optional', () => { const signature = AxSignature.fromZod( { @@ -339,7 +401,6 @@ describe('AxSignature.fromZod', () => { type: { name: 'date' }, }, ]); - }); it('warns on fallback by default and throws in strict mode', () => { @@ -371,7 +432,9 @@ describe('AxSignature.fromZod', () => { }, { strict: true } ) - ).toThrowError(/Unsupported Zod schema elements encountered during conversion/); + ).toThrowError( + /Unsupported Zod schema elements encountered during conversion/ + ); }); it('exposes conversion issues and debug helper output', () => { diff --git a/src/ax/dsp/zodToSignature.ts b/src/ax/dsp/zodToSignature.ts index 63c94cb6f..d4e7fccae 100644 --- a/src/ax/dsp/zodToSignature.ts +++ b/src/ax/dsp/zodToSignature.ts @@ -1,9 +1,7 @@ -import type { z, ZodObject, ZodTypeAny } from 'zod'; +import type { ZodTypeAny } from 'zod'; import type { AxField } from './sig.js'; -type ZodObjectLike = ZodObject; - type FieldTypeName = | 'string' | 'number' @@ -278,8 +276,7 @@ function hasStringCheck(def: ZodDef, keyword: string): boolean { ]; return candidates.some( (value) => - typeof value === 'string' && - value.toLowerCase().includes(lowerKeyword) + typeof value === 'string' && value.toLowerCase().includes(lowerKeyword) ); }); } @@ -346,10 +343,28 @@ function getFieldType( return { type: { name: 'json' } }; } if (typeof value === 'string' || typeof value === 'number') { + const literalOptions = [String(value)]; + if (context === 'input') { + recordIssue( + issues, + context, + path, + typeToken, + 'string', + 'Ax inputs do not support class fields; downgraded to string', + 'downgraded' + ); + return { + type: { + name: 'string', + options: literalOptions, + }, + }; + } return { type: { name: 'class', - options: [String(value)], + options: literalOptions, }, }; } @@ -474,17 +489,17 @@ function getFieldType( typeof literalValue !== 'number' && typeof literalValue !== 'boolean') ) { - recordIssue( - issues, - context, - path, - typeToken, - 'json', - 'Union includes non-stringifiable literal, falling back to json', - 'downgraded' - ); - return { type: { name: 'json' }, forceOptional }; - } + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Union includes literal that cannot be stringified, falling back to json', + 'downgraded' + ); + return { type: { name: 'json' }, forceOptional }; + } literalValues.push(literalValue as string | number | boolean); forceOptional = forceOptional || unwrapped.optional; continue; @@ -614,7 +629,7 @@ function getFieldType( } case 'record': { const keySchema = (schemaDef as { keyType?: ZodTypeAny }).keyType; - const valueSchema = (schemaDef as { valueType?: ZodTypeAny }).valueType; + const _valueSchema = (schemaDef as { valueType?: ZodTypeAny }).valueType; const keyToken = keySchema ? getTypeToken(keySchema) : undefined; const severity: ZodConversionSeverity = keyToken === 'string' ? 'downgraded' : 'unsupported'; @@ -622,15 +637,7 @@ function getFieldType( keyToken === 'string' ? 'Records with string keys remain json to preserve dynamic maps' : 'Records with non-string keys map to json'; - recordIssue( - issues, - context, - path, - typeToken, - 'json', - reason, - severity - ); + recordIssue(issues, context, path, typeToken, 'json', reason, severity); return { type: { name: 'json' } }; } case 'map': { @@ -642,15 +649,7 @@ function getFieldType( keyToken === 'string' ? 'Maps convert to json objects with dynamic keys' : 'Maps with non-string keys map to json'; - recordIssue( - issues, - context, - path, - typeToken, - 'json', - reason, - severity - ); + recordIssue(issues, context, path, typeToken, 'json', reason, severity); return { type: { name: 'json' } }; } case 'set': { @@ -660,7 +659,7 @@ function getFieldType( path, typeToken, 'json', - 'Sets are serialised as arrays; using json representation', + 'Sets are serialized as arrays; using json representation', 'downgraded' ); return { type: { name: 'json' } }; @@ -716,13 +715,21 @@ function getFieldType( } } -function getObjectShape(schema: ZodObjectLike): Record { +function getObjectShape(schema: ZodTypeAny): Record { if (typeof (schema as { shape?: unknown }).shape === 'function') { - return (schema as { shape: () => Record }).shape(); + return ( + schema as unknown as { + shape: () => Record; + } + ).shape(); } if ((schema as unknown as { shape?: Record }).shape) { - return (schema as unknown as { shape: Record }).shape; + return ( + schema as unknown as { + shape: Record; + } + ).shape; } return (schema as unknown as { _def: { shape: Record } }) @@ -730,10 +737,17 @@ function getObjectShape(schema: ZodObjectLike): Record { } export function zodObjectToSignatureFields( - schema: ZodObjectLike, + schema: ZodTypeAny, context: 'input' | 'output', options?: ZodObjectToSignatureOptions ): AxField[] { + const typeToken = getTypeToken(schema); + if (typeToken !== 'object') { + throw new TypeError( + `Expected a Zod object schema for ${context}, received ${typeToken ?? 'unknown'}` + ); + } + const shape = getObjectShape(schema); const fields: AxField[] = []; const basePath = options?.basePath ?? []; @@ -779,8 +793,22 @@ export function zodObjectToSignatureFields( return fields; } -export type InferZodInput = - T extends ZodObjectLike ? z.input : Record; - -export type InferZodOutput = - T extends ZodObjectLike ? z.output : Record; +type ZodInputType = T extends { _input: infer U } ? U : never; +type ZodOutputType = T extends { _output: infer U } ? U : never; +type NormalizeSchema = T extends { _input: any; _output: any } ? T : never; + +export type InferZodInput = NormalizeSchema extends infer S + ? S extends { _input: any; _output: any } + ? ZodInputType extends Record + ? ZodInputType + : Record + : Record + : never; + +export type InferZodOutput = NormalizeSchema extends infer S + ? S extends { _input: any; _output: any } + ? ZodOutputType extends Record + ? ZodOutputType + : Record + : Record + : never; diff --git a/src/ax/index.ts b/src/ax/index.ts index f751a27c6..e1dceae13 100644 --- a/src/ax/index.ts +++ b/src/ax/index.ts @@ -560,6 +560,7 @@ import { AxSignature, AxSignatureBuilder, type AxSignatureConfig, + type AxSignatureFromZodOptions, f, } from './dsp/sig.js'; import { AxStringUtil } from './dsp/strutil.js'; @@ -1252,6 +1253,7 @@ export type { AxRoutingResult }; export type { AxSamplePickerOptions }; export type { AxSetExamplesOptions }; export type { AxSignatureConfig }; +export type { AxSignatureFromZodOptions }; export type { AxSimpleClassifierForwardOptions }; export type { AxStreamingAssertion }; export type { AxStreamingEvent }; diff --git a/src/examples/zod-signature-example.ts b/src/examples/zod-signature-example.ts index 71a68d343..1fb80d3da 100644 --- a/src/examples/zod-signature-example.ts +++ b/src/examples/zod-signature-example.ts @@ -1,4 +1,4 @@ -import { AxSignature, ax } from '@ax-llm/ax'; +import { AxSignature } from '@ax-llm/ax'; import { z } from 'zod'; const bugReportInput = z.object({ @@ -35,11 +35,12 @@ const bugReportSignature = AxSignature.fromZod( const conversionIssues = bugReportSignature.getZodConversionIssues(); if (conversionIssues.length > 0) { - console.warn('[zod-signature-example] conversion issues stored on signature', conversionIssues); + console.warn( + '[zod-signature-example] conversion issues stored on signature', + conversionIssues + ); } -const triageAgent = ax(bugReportSignature); - async function main() { console.log('Bug report signature inputs:'); console.table(bugReportSignature.getInputFields()); From 54d08098a149a3eaf6f121d60da15f91568bdcea Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Thu, 16 Oct 2025 00:07:51 +0700 Subject: [PATCH 07/14] fix(signatures): prevent array element optionality leak --- src/ax/dsp/zodToSignature.test.ts | 24 ++++++++++++++++++++++++ src/ax/dsp/zodToSignature.ts | 6 ++++-- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/src/ax/dsp/zodToSignature.test.ts b/src/ax/dsp/zodToSignature.test.ts index 49f86bf94..7bd362bf5 100644 --- a/src/ax/dsp/zodToSignature.test.ts +++ b/src/ax/dsp/zodToSignature.test.ts @@ -310,6 +310,30 @@ describe('AxSignature.fromZod', () => { ]); }); + it('keeps array fields required when only the elements are optional', () => { + const signature = AxSignature.fromZod( + { + input: z.object({ + tags: z.array(z.string().optional()), + }), + output: z.object({ + ok: z.boolean(), + }), + }, + { warnOnFallback: false } + ); + + expect(signature.getInputFields()).toEqual([ + { + name: 'tags', + title: 'Tags', + type: { name: 'string', isArray: true }, + }, + ]); + + expect(signature.getZodConversionIssues()).toEqual([]); + }); + it('falls back to json for unsupported schema constructions', () => { const signature = AxSignature.fromZod( { diff --git a/src/ax/dsp/zodToSignature.ts b/src/ax/dsp/zodToSignature.ts index d4e7fccae..2f07a863b 100644 --- a/src/ax/dsp/zodToSignature.ts +++ b/src/ax/dsp/zodToSignature.ts @@ -596,13 +596,15 @@ function getFieldType( }; } - const result = { + // Array optionality should be controlled by wrappers on the array schema + // itself (handled via unwrapSchema). Optional element schemas must not + // leak up and mark the entire array field as optional. + const result: FieldTypeResult = { type: { name: elementType.type.name, isArray: true, options: elementType.type.options, }, - forceOptional: elementSchema.optional || elementType.forceOptional, }; if (context === 'input' && result.type.name === 'class') { From 36ddcc7d01cffbc396c1ff4540bff16a248cb9d3 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Thu, 16 Oct 2025 00:09:56 +0700 Subject: [PATCH 08/14] docs(signatures): clarify zod array optionality --- docs/SIGNATURES.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/SIGNATURES.md b/docs/SIGNATURES.md index 999d9e545..93d3632f4 100644 --- a/docs/SIGNATURES.md +++ b/docs/SIGNATURES.md @@ -224,7 +224,7 @@ const signature = AxSignature.fromZod({ **Mapping highlights** - `z.string()`, `z.number()`, `z.boolean()`, `z.date()` map to matching Ax field types. String refinements such as `.url()` and `.datetime()` become `url`/`datetime` field types. -- `z.array()` becomes Ax arrays; nested arrays fall back to `json`. +- `z.array()` becomes Ax arrays; nested arrays fall back to `json`. Optionality flows from wrappers on the array schema (`z.array(...).optional()`), while optional elements (`z.array(z.string().optional())`) still produce required arrays at runtime. - Literal unions/`z.enum`/`z.nativeEnum` values become classifications; input unions are exposed as `string` fields with `options` metadata (Ax intentionally disallows `class` inputs). - Optional/nullable/default/catch wrappers automatically mark the field optional. - Records, maps, discriminated unions, and other dynamic structures stay as `json`, but Ax marks them as **downgraded** so you can adjust the schema (or rely on `strict: true`). @@ -259,7 +259,8 @@ signature.reportZodConversionIssues(); Need a quick readout before wiring it in? Call `AxSignature.debugZodConversion({ input, output })` to get both the signature -and a ready-made downgrade report. +and a ready-made downgrade report. You can pass `{ logger: yourTelemetry }` to +pipe the downgrade issues into custom observability tooling. **Standard Schema?** From 214fe9d07f7a52f7666a45c1040d45278a1df107 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Thu, 16 Oct 2025 00:20:15 +0700 Subject: [PATCH 09/14] docs(examples): add zod flow integration demo --- src/examples/zod-flow-example.ts | 95 ++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 src/examples/zod-flow-example.ts diff --git a/src/examples/zod-flow-example.ts b/src/examples/zod-flow-example.ts new file mode 100644 index 000000000..21f4d132a --- /dev/null +++ b/src/examples/zod-flow-example.ts @@ -0,0 +1,95 @@ +import { AxFlow, AxMockAIService, AxSignature } from '@ax-llm/ax'; +import { z } from 'zod'; + +const ticketInputSchema = z.object({ + summary: z + .string() + .min(1) + .describe('Short description of the incident being triaged'), + details: z + .string() + .min(10) + .describe('Full report including any reproduction steps'), + severity: z.enum(['low', 'medium', 'high']), + tags: z + .array(z.string().optional()) + .describe('Optional labels shared with observability dashboards'), +}); + +const ticketOutputSchema = z.object({ + recommendedAction: z.enum(['escalate', 'monitor', 'close']), + priority: z.enum(['P0', 'P1', 'P2', 'P3']), + suggestedOwners: z.array(z.string()).optional(), +}); + +const { signature: ticketSignature, issues } = AxSignature.debugZodConversion( + { + description: 'Classify incoming incidents and propose a triage plan', + input: ticketInputSchema, + output: ticketOutputSchema, + }, + { + logger: (conversionIssues) => { + if (conversionIssues.length === 0) { + console.info('[zod-flow-example] no downgrades detected'); + } else { + console.warn( + '[zod-flow-example] zod conversion issues', + conversionIssues + ); + } + }, + } +); + +type TicketInput = z.input; +type TicketOutput = z.output; + +const triageFlow = AxFlow.create() + .node('triage', ticketSignature) + .execute('triage', (state) => state) + .map((state) => ({ + recommendedAction: state.triageResult.recommendedAction, + priority: state.triageResult.priority, + suggestedOwners: state.triageResult.suggestedOwners, + })); + +const mockAI = new AxMockAIService({ + chatResponse: async () => ({ + results: [ + { + index: 0, + content: + 'Recommended Action: escalate\nPriority: P1\nSuggested Owners: ["incident-response","platform-oncall"]', + finishReason: 'stop', + }, + ], + modelUsage: { + ai: 'mock', + model: 'zod-flow-mock', + tokens: { promptTokens: 42, completionTokens: 58, totalTokens: 100 }, + }, + }), +}); + +const sampleTicket: TicketInput = { + summary: 'Spike in 500s on checkout service', + details: + 'Multiple customers report failures when completing checkout. Error rate jumped to 7% in the last 10 minutes.', + severity: 'high', + tags: ['checkout', 'payments'], +}; + +async function main() { + console.log('[zod-flow-example] ticket signature:', ticketSignature.toString()); + if (issues.length > 0) { + console.warn('[zod-flow-example] stored issues:', issues); + } + + const result = await triageFlow.forward(mockAI, sampleTicket); + + console.log('[zod-flow-example] triage result:'); + console.dir(result, { depth: null }); +} + +void main(); From bd5f002d174e91d56f98bf79870e3b8e850d594b Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Thu, 16 Oct 2025 00:22:53 +0700 Subject: [PATCH 10/14] docs(examples): switch zod flow demo to factory --- src/examples/zod-flow-example.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/examples/zod-flow-example.ts b/src/examples/zod-flow-example.ts index 21f4d132a..28c328ae3 100644 --- a/src/examples/zod-flow-example.ts +++ b/src/examples/zod-flow-example.ts @@ -1,4 +1,4 @@ -import { AxFlow, AxMockAIService, AxSignature } from '@ax-llm/ax'; +import { AxMockAIService, AxSignature, flow } from '@ax-llm/ax'; import { z } from 'zod'; const ticketInputSchema = z.object({ @@ -45,7 +45,7 @@ const { signature: ticketSignature, issues } = AxSignature.debugZodConversion( type TicketInput = z.input; type TicketOutput = z.output; -const triageFlow = AxFlow.create() +const triageFlow = flow() .node('triage', ticketSignature) .execute('triage', (state) => state) .map((state) => ({ From 327672ad53751f1fba81b23edf1092dea68e7c2a Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Thu, 16 Oct 2025 00:32:18 +0700 Subject: [PATCH 11/14] docs(examples): tighten zod flow mock typing --- src/examples/zod-flow-example.ts | 48 +++++++++++++++++++++----------- 1 file changed, 31 insertions(+), 17 deletions(-) diff --git a/src/examples/zod-flow-example.ts b/src/examples/zod-flow-example.ts index 28c328ae3..c98e21921 100644 --- a/src/examples/zod-flow-example.ts +++ b/src/examples/zod-flow-example.ts @@ -1,4 +1,9 @@ -import { AxMockAIService, AxSignature, flow } from '@ax-llm/ax'; +import { + AxMockAIService, + AxSignature, + flow, + type AxChatResponse, +} from '@ax-llm/ax'; import { z } from 'zod'; const ticketInputSchema = z.object({ @@ -54,22 +59,28 @@ const triageFlow = flow() suggestedOwners: state.triageResult.suggestedOwners, })); -const mockAI = new AxMockAIService({ - chatResponse: async () => ({ - results: [ - { - index: 0, - content: - 'Recommended Action: escalate\nPriority: P1\nSuggested Owners: ["incident-response","platform-oncall"]', - finishReason: 'stop', - }, - ], - modelUsage: { - ai: 'mock', - model: 'zod-flow-mock', - tokens: { promptTokens: 42, completionTokens: 58, totalTokens: 100 }, +const mockResponse: AxChatResponse = { + results: [ + { + index: 0, + content: + 'Recommended Action: escalate\nPriority: P1\nSuggested Owners: ["incident-response","platform-oncall"]', + finishReason: 'stop', }, - }), + ], + modelUsage: { + ai: 'mock', + model: 'zod-flow-mock', + tokens: { + promptTokens: 42, + completionTokens: 58, + totalTokens: 100, + }, + }, +}; + +const mockAI = new AxMockAIService({ + chatResponse: async () => mockResponse, }); const sampleTicket: TicketInput = { @@ -81,7 +92,10 @@ const sampleTicket: TicketInput = { }; async function main() { - console.log('[zod-flow-example] ticket signature:', ticketSignature.toString()); + console.log( + '[zod-flow-example] ticket signature:', + ticketSignature.toString() + ); if (issues.length > 0) { console.warn('[zod-flow-example] stored issues:', issues); } From 0f4f9242682c5f769bef42ddb22b3db70470d298 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Thu, 16 Oct 2025 00:44:53 +0700 Subject: [PATCH 12/14] docs(examples): add zod live openai demo --- src/examples/zod-flow-example.ts | 109 --------------------------- src/examples/zod-live-example.ts | 125 +++++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 109 deletions(-) delete mode 100644 src/examples/zod-flow-example.ts create mode 100644 src/examples/zod-live-example.ts diff --git a/src/examples/zod-flow-example.ts b/src/examples/zod-flow-example.ts deleted file mode 100644 index c98e21921..000000000 --- a/src/examples/zod-flow-example.ts +++ /dev/null @@ -1,109 +0,0 @@ -import { - AxMockAIService, - AxSignature, - flow, - type AxChatResponse, -} from '@ax-llm/ax'; -import { z } from 'zod'; - -const ticketInputSchema = z.object({ - summary: z - .string() - .min(1) - .describe('Short description of the incident being triaged'), - details: z - .string() - .min(10) - .describe('Full report including any reproduction steps'), - severity: z.enum(['low', 'medium', 'high']), - tags: z - .array(z.string().optional()) - .describe('Optional labels shared with observability dashboards'), -}); - -const ticketOutputSchema = z.object({ - recommendedAction: z.enum(['escalate', 'monitor', 'close']), - priority: z.enum(['P0', 'P1', 'P2', 'P3']), - suggestedOwners: z.array(z.string()).optional(), -}); - -const { signature: ticketSignature, issues } = AxSignature.debugZodConversion( - { - description: 'Classify incoming incidents and propose a triage plan', - input: ticketInputSchema, - output: ticketOutputSchema, - }, - { - logger: (conversionIssues) => { - if (conversionIssues.length === 0) { - console.info('[zod-flow-example] no downgrades detected'); - } else { - console.warn( - '[zod-flow-example] zod conversion issues', - conversionIssues - ); - } - }, - } -); - -type TicketInput = z.input; -type TicketOutput = z.output; - -const triageFlow = flow() - .node('triage', ticketSignature) - .execute('triage', (state) => state) - .map((state) => ({ - recommendedAction: state.triageResult.recommendedAction, - priority: state.triageResult.priority, - suggestedOwners: state.triageResult.suggestedOwners, - })); - -const mockResponse: AxChatResponse = { - results: [ - { - index: 0, - content: - 'Recommended Action: escalate\nPriority: P1\nSuggested Owners: ["incident-response","platform-oncall"]', - finishReason: 'stop', - }, - ], - modelUsage: { - ai: 'mock', - model: 'zod-flow-mock', - tokens: { - promptTokens: 42, - completionTokens: 58, - totalTokens: 100, - }, - }, -}; - -const mockAI = new AxMockAIService({ - chatResponse: async () => mockResponse, -}); - -const sampleTicket: TicketInput = { - summary: 'Spike in 500s on checkout service', - details: - 'Multiple customers report failures when completing checkout. Error rate jumped to 7% in the last 10 minutes.', - severity: 'high', - tags: ['checkout', 'payments'], -}; - -async function main() { - console.log( - '[zod-flow-example] ticket signature:', - ticketSignature.toString() - ); - if (issues.length > 0) { - console.warn('[zod-flow-example] stored issues:', issues); - } - - const result = await triageFlow.forward(mockAI, sampleTicket); - - console.log('[zod-flow-example] triage result:'); - console.dir(result, { depth: null }); -} - -void main(); diff --git a/src/examples/zod-live-example.ts b/src/examples/zod-live-example.ts new file mode 100644 index 000000000..0fbbee445 --- /dev/null +++ b/src/examples/zod-live-example.ts @@ -0,0 +1,125 @@ +import { AxAI, AxAIOpenAIModel, AxGen, AxSignature, flow } from '@ax-llm/ax'; +import { z } from 'zod'; + +const incidentInputSchema = z.object({ + summary: z + .string() + .min(1) + .describe('Short description of the incident being triaged'), + details: z + .string() + .min(10) + .describe('Full report including any reproduction steps'), + severity: z.enum(['low', 'medium', 'high']), + tags: z + .array(z.string()) + .optional() + .describe('Optional labels shared with observability dashboards'), +}); + +const incidentOutputSchema = z.object({ + recommendedAction: z.enum(['escalate', 'monitor', 'close']), + priority: z.enum(['P0', 'P1', 'P2', 'P3']), + suggestedOwners: z.array(z.string()).optional(), +}); + +const { signature: incidentSignature, issues } = AxSignature.debugZodConversion( + { + description: 'Classify incoming incidents and propose a triage plan', + input: incidentInputSchema, + output: incidentOutputSchema, + }, + { + logger: (conversionIssues) => { + if (conversionIssues.length === 0) { + console.info('[zod-live-example] no downgrades detected'); + } else { + console.warn('[zod-live-example] conversion issues', conversionIssues); + } + }, + } +); + +type IncidentInput = z.input; +type IncidentOutput = z.output; + +const sampleIncident: IncidentInput = { + summary: 'Spike in 500 errors on checkout service', + details: + 'Multiple customers report failures when completing checkout. Error rate jumped to 7% in the last 10 minutes.', + severity: 'high', + tags: ['checkout', 'payments'], +}; + +function createOpenAIAI(): AxAI { + const apiKey = process.env.OPENAI_APIKEY; + + if (!apiKey) { + throw new Error( + 'Set OPENAI_APIKEY before running this example (export OPENAI_APIKEY=...)' + ); + } + + return new AxAI({ + name: 'openai', + apiKey, + config: { + model: AxAIOpenAIModel.GPT5Mini, + maxTokens: 512, + }, + }); +} + +async function runDirectExample(ai: AxAI): Promise { + const generator = new AxGen( + incidentSignature, + { + description: + 'Triages incidents into actionable recommendations using Zod-derived signature', + } + ); + + console.log('\n[zod-live-example] running AxGen.fromZod direct example...'); + + const directResult = await generator.forward(ai, sampleIncident); + + console.log('[zod-live-example] direct result:'); + console.dir(directResult, { depth: null }); +} + +async function runFlowExample(ai: AxAI): Promise { + const triageFlow = flow() + .node('triage', incidentSignature) + .execute('triage', (state) => state) + .map((state) => ({ + recommendedAction: state.triageResult.recommendedAction, + priority: state.triageResult.priority, + suggestedOwners: state.triageResult.suggestedOwners, + })); + + console.log('\n[zod-live-example] running AxFlow + fromZod example...'); + + const flowResult = await triageFlow.forward(ai, sampleIncident); + + console.log('[zod-live-example] flow result:'); + console.dir(flowResult, { depth: null }); +} + +async function main(): Promise { + console.log('[zod-live-example] ticket signature:'); + console.log(incidentSignature.toString()); + + if (issues.length > 0) { + console.warn('[zod-live-example] stored issues:', issues); + } + + try { + const ai = createOpenAIAI(); + await runDirectExample(ai); + await runFlowExample(ai); + } catch (error) { + console.error('[zod-live-example] failed to run example:', error); + } +} + +void main(); From bade7b5491b4f01591688c9104c8b0d611c6bcf5 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Thu, 16 Oct 2025 01:23:00 +0700 Subject: [PATCH 13/14] fix(signatures): flag zod effects and pipes --- src/ax/dsp/zodToSignature.test.ts | 70 +++++++++++++++++++++++++++++++ src/ax/dsp/zodToSignature.ts | 41 ++++++++++-------- 2 files changed, 93 insertions(+), 18 deletions(-) diff --git a/src/ax/dsp/zodToSignature.test.ts b/src/ax/dsp/zodToSignature.test.ts index 7bd362bf5..d9b948020 100644 --- a/src/ax/dsp/zodToSignature.test.ts +++ b/src/ax/dsp/zodToSignature.test.ts @@ -111,6 +111,18 @@ describe('AxSignature.fromZod', () => { const inputFields = signature.getInputFields(); + if (process.env.DEBUG_ZOD_TRANSFORMS) { + // eslint-disable-next-line no-console + console.log('nullable/default input', inputFields); + // eslint-disable-next-line no-console + console.log('nullable/default output', outputFields); + // eslint-disable-next-line no-console + console.log( + 'nullable/default issues', + signature.getZodConversionIssues() + ); + } + expect(inputFields).toEqual([ { name: 'tone', @@ -265,6 +277,64 @@ describe('AxSignature.fromZod', () => { type: { name: 'string' }, }, ]); + + expect(signature.getZodConversionIssues()).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + path: ['input', 'pipeline'], + fallbackType: 'json', + severity: 'downgraded', + }), + ]) + ); + }); + + it('downgrades Zod transforms to json fields with issues', () => { + const signature = AxSignature.fromZod( + { + input: z.object({ + transformed: z.string().transform((value) => value.length), + }), + output: z.object({ + converted: z + .string() + .transform((value) => ({ length: value.length })), + }), + }, + { warnOnFallback: false } + ); + + expect(signature.getInputFields()).toEqual([ + { + name: 'transformed', + title: 'Transformed', + type: { name: 'json' }, + }, + ]); + + expect(signature.getOutputFields()).toEqual([ + { + name: 'converted', + title: 'Converted', + type: { name: 'json' }, + }, + ]); + + expect(signature.getZodConversionIssues()).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + path: ['input', 'transformed'], + fallbackType: 'json', + severity: 'downgraded', + }), + expect.objectContaining({ + path: ['output', 'converted'], + fallbackType: 'json', + severity: 'downgraded', + schemaType: expect.any(String), + }), + ]) + ); }); it('handles arrays and nested objects by falling back to json where needed', () => { diff --git a/src/ax/dsp/zodToSignature.ts b/src/ax/dsp/zodToSignature.ts index 2f07a863b..019875c03 100644 --- a/src/ax/dsp/zodToSignature.ts +++ b/src/ax/dsp/zodToSignature.ts @@ -157,24 +157,15 @@ function unwrapSchema(schema: ZodTypeAny): UnwrappedSchema { continue; } - if (typeToken === 'effects') { - const next = getDef(current).schema; - if (!next || typeof next !== 'object') { - break; - } - current = next as ZodTypeAny; - description ??= current.description; - continue; - } - - if (typeToken === 'pipeline') { - const next = getDef(current).out; - if (!next || typeof next !== 'object') { - break; - } - current = next as ZodTypeAny; - description ??= current.description; - continue; + if ( + typeToken === 'effects' || + typeToken === 'pipeline' || + typeToken === 'pipe' + ) { + // Zod effects/pipeline nodes wrap transformations/refinements that can + // change runtime data. We keep the wrapper intact so the caller can + // detect the downgrade (handled in getFieldType). + break; } if (typeToken === 'branded') { @@ -666,6 +657,20 @@ function getFieldType( ); return { type: { name: 'json' } }; } + case 'effects': + case 'pipeline': + case 'pipe': { + recordIssue( + issues, + context, + path, + typeToken, + 'json', + 'Zod transformations (effects/pipe) may alter runtime types; falling back to json', + 'downgraded' + ); + return { type: { name: 'json' } }; + } case 'discriminatedUnion': { recordIssue( issues, From c73c75e718a8d2b1452f6b6251882f881e6a1253 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Thu, 16 Oct 2025 10:41:23 +0700 Subject: [PATCH 14/14] feat(signatures): add signature to zod conversion --- README.md | 10 ++ docs/SIGNATURES.md | 29 +++++ src/ax/dsp/sig.ts | 88 ++++++++++++++ src/ax/dsp/signatureToZod.test.ts | 155 +++++++++++++++++++++++++ src/ax/dsp/signatureToZod.ts | 161 ++++++++++++++++++++++++++ src/ax/index.test-d.ts | 9 ++ src/ax/index.ts | 2 + src/examples/zod-signature-example.ts | 17 +++ 8 files changed, 471 insertions(+) create mode 100644 src/ax/dsp/signatureToZod.test.ts create mode 100644 src/ax/dsp/signatureToZod.ts diff --git a/README.md b/README.md index 814b4e786..2677ccc96 100644 --- a/README.md +++ b/README.md @@ -193,6 +193,15 @@ AxSignature.debugZodConversion({ // Emit a warning-style report when issues exist ticketSignature.reportZodConversionIssues(); +// Need the Zod schemas back out (e.g. for adapters)? +const { input: inputSchema, output: outputSchema } = ticketSignature.toZod({ + warnOnFallback: false, +}); +if (inputSchema && outputSchema) { + type TicketInput = z.input; + type TicketOutput = z.output; +} + const summarize = ax(ticketSignature); ``` @@ -203,6 +212,7 @@ const summarize = ax(ticketSignature); - ✅ **Type-Safe Everything** - Full TypeScript support with auto-completion - ✅ **Streaming First** - Real-time responses with validation - ✅ **Zod-Friendly** - Convert schemas with automatic fallbacks and warnings +- ✅ **Round-Trip Friendly** - Regenerate Zod objects from signatures when you need adapters - ✅ **Downgrade Awareness** - Records/maps/unions stay `json` but are flagged so you can adjust early - ✅ **Multi-Modal** - Images, audio, text in the same signature - ✅ **Smart Optimization** - Automatic prompt tuning with MiPRO diff --git a/docs/SIGNATURES.md b/docs/SIGNATURES.md index 93d3632f4..59a6d067b 100644 --- a/docs/SIGNATURES.md +++ b/docs/SIGNATURES.md @@ -262,6 +262,35 @@ Need a quick readout before wiring it in? Call and a ready-made downgrade report. You can pass `{ logger: yourTelemetry }` to pipe the downgrade issues into custom observability tooling. +**Round-trip back to Zod** + +Already have a signature and need Zod validators (for adapters, loaders, or +form tooling)? Call `signature.toZod()` to reconstruct the schemas: + +```typescript +import { z } from 'zod'; + +const { input: zodInput, output: zodOutput, issues } = signature.toZod({ + warnOnFallback: false, +}); + +if (zodInput) { + type InputPayload = z.input; +} +if (zodOutput) { + type OutputPayload = z.output; +} + +if (issues.length > 0) { + console.warn('Some fields were widened for Zod conversion', issues); +} +``` + +Most Ax field types map cleanly (`string`, `number`, `boolean`, `image`, `file`, +etc.). We fall back to permissive schemas such as `z.any()` for types without a +direct Zod equivalent and report them through the `issues` array (or throw when +`strict: true`). + **Standard Schema?** Standard Schema libraries (Effect Schema, Valibot, ArkType, etc.) publish adapters to Zod or JSON Schema. Convert with your preferred tool (for example [`xsschema`](https://xsai.js.org/docs/packages-top/xsschema)) and feed the resulting Zod schema into `AxSignature.fromZod`. diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts index 669994976..ab16d628e 100644 --- a/src/ax/dsp/sig.ts +++ b/src/ax/dsp/sig.ts @@ -17,6 +17,11 @@ import { zodObjectToSignatureFields, } from './zodToSignature.js'; export type { ZodConversionIssue } from './zodToSignature.js'; +import { + signatureFieldsToZodObject, + type SignatureToZodIssue, +} from './signatureToZod.js'; +export type { SignatureToZodIssue } from './signatureToZod.js'; // Interface for programmatically defining field types export interface AxFieldType { readonly type: @@ -574,6 +579,21 @@ export type AxSignatureFromZodOptions = { readonly onIssues?: (issues: readonly ZodConversionIssue[]) => void; }; +export type AxSignatureToZodOptions = { + /** + * When true, throw if any field must fall back to a permissive Zod schema. + */ + readonly strict?: boolean; + /** + * When true (default), emit a console warning when downgrade issues occur. + */ + readonly warnOnFallback?: boolean; + /** + * Receives downgrade metadata for custom handling. + */ + readonly onIssues?: (issues: readonly SignatureToZodIssue[]) => void; +}; + export class AxSignature< _TInput extends Record = Record, _TOutput extends Record = Record, @@ -811,6 +831,30 @@ export class AxSignature< console.warn(message); } + private static formatSignatureToZodIssues( + issues: ReadonlyArray + ): string { + const details = issues.map((issue) => { + const target = issue.path.join('.'); + return `[${issue.context}] ${target} → ${issue.fallback} (${issue.severity}: ${issue.reason})`; + }); + return `Review the following conversions:\n${details.join('\n')}`; + } + + private static warnAboutSignatureToZodIssues( + issues: ReadonlyArray + ): void { + const messageLines = issues.map((issue) => { + const target = issue.path.join('.'); + return ` - [${issue.context}] ${target} → ${issue.fallback}: ${issue.reason}`; + }); + const message = [ + '[AxSignature.toZod] Some fields were downgraded to permissive Zod schemas.', + ...messageLines, + ].join('\n'); + console.warn(message); + } + public static debugZodConversion( config: { description?: string; @@ -861,6 +905,50 @@ export class AxSignature< }; } + public toZod(options?: AxSignatureToZodOptions): { + readonly input?: ReturnType; + readonly output?: ReturnType; + readonly issues: readonly SignatureToZodIssue[]; + } { + const issues: SignatureToZodIssue[] = []; + + const input = this.inputFields.length + ? signatureFieldsToZodObject(this.inputFields, 'input', { + issues, + basePath: ['input'], + }) + : undefined; + + const output = this.outputFields.length + ? signatureFieldsToZodObject(this.outputFields, 'output', { + issues, + basePath: ['output'], + }) + : undefined; + + if (issues.length > 0) { + options?.onIssues?.(issues); + + if (options?.strict) { + throw new AxSignatureValidationError( + 'Unsupported Ax field types encountered during Zod conversion', + undefined, + AxSignature.formatSignatureToZodIssues(issues) + ); + } + + if (options?.warnOnFallback !== false) { + AxSignature.warnAboutSignatureToZodIssues(issues); + } + } + + return { + input, + output, + issues: Object.freeze([...issues]), + }; + } + public getZodConversionIssues = (): readonly ZodConversionIssue[] => this.zodConversionIssues; diff --git a/src/ax/dsp/signatureToZod.test.ts b/src/ax/dsp/signatureToZod.test.ts new file mode 100644 index 000000000..f087f1d3b --- /dev/null +++ b/src/ax/dsp/signatureToZod.test.ts @@ -0,0 +1,155 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { AxSignature } from './sig.js'; + +describe('AxSignature.toZod', () => { + it('converts primitive fields with optional arrays to Zod schemas', () => { + const signature = AxSignature.create( + 'query: string, tags?: string[] -> responseText: string' + ); + + const { input, output, issues } = signature.toZod({ + warnOnFallback: false, + }); + + expect(issues).toEqual([]); + expect(input).toBeDefined(); + expect(output).toBeDefined(); + + expect(input?.parse({ query: 'hello world' })).toEqual({ + query: 'hello world', + }); + expect( + output?.parse({ + responseText: 'done', + }) + ).toEqual({ responseText: 'done' }); + expect(() => input?.parse({ tags: [] })).toThrowError(); + }); + + it('preserves class options via enums in the Zod schema', () => { + const signature = new AxSignature({ + inputs: [ + { + name: 'userInput', + type: { name: 'string' }, + }, + ], + outputs: [ + { + name: 'status', + type: { name: 'class', options: ['ok', 'error'] }, + }, + ], + }); + + const { output, issues } = signature.toZod({ warnOnFallback: false }); + expect(issues).toEqual([]); + expect(output?.parse({ status: 'ok' })).toEqual({ status: 'ok' }); + expect(() => output?.parse({ status: 'unknown' })).toThrowError(); + }); + + it('emits downgrade issues for unsupported field types and respects strict mode', () => { + const signature = new AxSignature({ + inputs: [ + { + name: 'userInput', + type: { name: 'string' }, + }, + ], + outputs: [ + { + name: 'customResult', + type: { name: 'mystery' }, + }, + ], + }); + + const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const result = signature.toZod(); + + expect(result.output).toBeDefined(); + expect(result.issues).toEqual([ + expect.objectContaining({ + context: 'output', + path: ['output', 'customResult'], + fallback: 'z.any()', + severity: 'unsupported', + }), + ]); + expect(warnSpy).toHaveBeenCalled(); + warnSpy.mockRestore(); + + expect(() => + signature.toZod({ warnOnFallback: false, strict: true }) + ).toThrowError( + /Unsupported Ax field types encountered during Zod conversion/ + ); + }); + + it('maps multimodal fields to structured Zod objects', () => { + const signature = new AxSignature({ + inputs: [ + { + name: 'incidentImage', + type: { name: 'image' }, + }, + { + name: 'incidentAttachment', + type: { name: 'file' }, + isOptional: true, + }, + { + name: 'referenceUrl', + type: { name: 'url' }, + isOptional: true, + }, + ], + outputs: [ + { + name: 'summaryText', + type: { name: 'string' }, + }, + ], + }); + + const { input, output, issues } = signature.toZod({ + warnOnFallback: false, + }); + + expect(issues).toEqual([]); + const parsedInput = input?.parse({ + incidentImage: { mimeType: 'image/png', data: 'base64' }, + incidentAttachment: { + mimeType: 'application/pdf', + fileUri: 's3://file.pdf', + }, + referenceUrl: { url: 'https://example.com/details', title: 'Details' }, + }); + expect(parsedInput).toEqual({ + incidentImage: { mimeType: 'image/png', data: 'base64' }, + incidentAttachment: { + mimeType: 'application/pdf', + fileUri: 's3://file.pdf', + }, + referenceUrl: { url: 'https://example.com/details', title: 'Details' }, + }); + + expect( + input?.parse({ + incidentImage: { mimeType: 'image/png', data: 'base64' }, + referenceUrl: 'https://example.com/details', + }) + ).toEqual({ + incidentImage: { mimeType: 'image/png', data: 'base64' }, + referenceUrl: 'https://example.com/details', + }); + + const parsedOutput = output?.parse({ + summaryText: 'Incident summary created', + }); + expect(parsedOutput).toEqual({ + summaryText: 'Incident summary created', + }); + }); +}); diff --git a/src/ax/dsp/signatureToZod.ts b/src/ax/dsp/signatureToZod.ts new file mode 100644 index 000000000..b919d46e0 --- /dev/null +++ b/src/ax/dsp/signatureToZod.ts @@ -0,0 +1,161 @@ +import { z, type ZodTypeAny } from 'zod'; + +import type { AxField } from './sig.js'; + +export type SignatureToZodSeverity = 'downgraded' | 'unsupported'; + +export type SignatureToZodIssue = { + readonly context: 'input' | 'output'; + readonly path: readonly string[]; + readonly fieldType: string; + readonly fallback: string; + readonly reason: string; + readonly severity: SignatureToZodSeverity; +}; + +export type SignatureFieldsToZodOptions = { + readonly issues?: SignatureToZodIssue[]; + readonly basePath?: readonly string[]; +}; + +const UNIQUE_URL_SCHEMA = z.object({ + url: z.string().url(), + title: z.string().optional(), + description: z.string().optional(), +}); + +function recordIssue( + issues: SignatureToZodIssue[] | undefined, + context: 'input' | 'output', + path: readonly string[], + fieldType: string, + fallback: string, + reason: string, + severity: SignatureToZodSeverity +): void { + if (!issues) { + return; + } + issues.push({ + context, + fallback, + fieldType, + path: [...path], + reason, + severity, + }); +} + +function createEnum(options: readonly string[]) { + const unique = Array.from(new Set(options)); + if (unique.length === 0) { + return undefined; + } + const [first, ...rest] = unique; + return z.enum([first, ...rest] as [string, ...string[]]); +} + +function buildBaseSchema( + field: Readonly, + context: 'input' | 'output', + issues: SignatureToZodIssue[] | undefined, + path: readonly string[] +): ZodTypeAny { + const fieldType = field.type?.name ?? 'string'; + switch (fieldType) { + case 'string': { + if (field.type?.options?.length) { + const enumSchema = createEnum(field.type.options); + if (enumSchema) { + return enumSchema; + } + } + return z.string(); + } + case 'number': + return z.number(); + case 'boolean': + return z.boolean(); + case 'json': + return z.any(); + case 'image': + return z.object({ + mimeType: z.string(), + data: z.string(), + }); + case 'file': + return z.union([ + z.object({ + mimeType: z.string(), + data: z.string(), + }), + z.object({ + mimeType: z.string(), + fileUri: z.string(), + }), + ]); + case 'url': + return z.union([z.string().url(), UNIQUE_URL_SCHEMA]); + case 'date': + return z.union([z.date(), z.string()]); + case 'datetime': + return z.union([z.date(), z.string()]); + case 'class': { + if (field.type?.options?.length) { + const enumSchema = createEnum(field.type.options); + if (enumSchema) { + return enumSchema; + } + } + recordIssue( + issues, + context, + path, + fieldType, + 'z.string()', + 'Class field without options maps to string', + 'downgraded' + ); + return z.string(); + } + case 'code': + return z.string(); + default: + recordIssue( + issues, + context, + path, + fieldType, + 'z.any()', + 'Field type not directly representable in Zod, using z.any()', + 'unsupported' + ); + return z.any(); + } +} + +export function signatureFieldsToZodObject( + fields: readonly AxField[], + context: 'input' | 'output', + options?: SignatureFieldsToZodOptions +) { + const basePath = options?.basePath ?? []; + const shape: Record = {}; + + for (const field of fields) { + const path = [...basePath, field.name]; + let schema = buildBaseSchema(field, context, options?.issues, path); + if (field.type?.isArray) { + schema = z.array(schema); + } + if (field.isOptional) { + schema = schema.optional(); + } + if (field.description) { + schema = schema.describe(field.description); + } + shape[field.name] = schema; + } + + return z.object(shape); +} diff --git a/src/ax/index.test-d.ts b/src/ax/index.test-d.ts index e61561653..260a99a44 100644 --- a/src/ax/index.test-d.ts +++ b/src/ax/index.test-d.ts @@ -4,6 +4,7 @@ import { z } from 'zod'; // === Typesafe Signature Tests === import { AxSignature } from './dsp/sig.js'; +import type { SignatureToZodIssue } from './dsp/sig.js'; import type { ZodConversionIssue } from './dsp/zodToSignature.js'; // Test basic signature type inference @@ -96,6 +97,14 @@ import type { AxExamples } from './dsp/types.js'; const testSig = AxSignature.create('userInput: string -> responseText: string'); +const toZodResult = testSig.toZod({ + onIssues: (issues) => { + expectType(issues); + }, + warnOnFallback: false, +}); +expectType(toZodResult.issues); + // Test appendInputField type inference const withAppendedInput = testSig.appendInputField('contextInfo', { type: 'string', diff --git a/src/ax/index.ts b/src/ax/index.ts index e1dceae13..87b3bf621 100644 --- a/src/ax/index.ts +++ b/src/ax/index.ts @@ -561,6 +561,7 @@ import { AxSignatureBuilder, type AxSignatureConfig, type AxSignatureFromZodOptions, + type AxSignatureToZodOptions, f, } from './dsp/sig.js'; import { AxStringUtil } from './dsp/strutil.js'; @@ -1254,6 +1255,7 @@ export type { AxSamplePickerOptions }; export type { AxSetExamplesOptions }; export type { AxSignatureConfig }; export type { AxSignatureFromZodOptions }; +export type { AxSignatureToZodOptions }; export type { AxSimpleClassifierForwardOptions }; export type { AxStreamingAssertion }; export type { AxStreamingEvent }; diff --git a/src/examples/zod-signature-example.ts b/src/examples/zod-signature-example.ts index 1fb80d3da..4c30670db 100644 --- a/src/examples/zod-signature-example.ts +++ b/src/examples/zod-signature-example.ts @@ -48,6 +48,23 @@ async function main() { console.log('\nBug report signature outputs:'); console.table(bugReportSignature.getOutputFields()); + const { + input: zodInputSchema, + output: zodOutputSchema, + issues, + } = bugReportSignature.toZod({ warnOnFallback: false }); + + if (issues.length > 0) { + console.warn('[zod-signature-example] signature → Zod issues', issues); + } + + if (zodInputSchema && zodOutputSchema) { + console.log('\nRound-tripped Zod schemas:', { + inputShape: Object.keys(zodInputSchema.shape), + outputShape: Object.keys(zodOutputSchema.shape), + }); + } + // In a real flow you would pass an AxAI instance here. // This example just shows the structure defined by the schema. const fakeReport = {