From 33a26cfca9927186d741a49b0395d1a28bbd5442 Mon Sep 17 00:00:00 2001 From: Tom Date: Sat, 18 Oct 2025 18:20:29 +0700 Subject: [PATCH 1/3] Add OpenAI Zod generation example --- docs/ZOD_INTEGRATION.md | 200 ++++++++++ package-lock.json | 3 +- src/ax/dsp/generate.test.ts | 138 +++++++ src/ax/dsp/generate.ts | 37 +- src/ax/dsp/processResponse.ts | 4 +- src/ax/dsp/program.ts | 36 +- src/ax/dsp/sig.test.ts | 98 ++++- src/ax/dsp/sig.ts | 78 ++++ src/ax/dsp/types.ts | 2 + src/ax/index.ts | 16 + src/ax/package.json | 3 +- src/ax/zod/assertion.ts | 62 +++ src/ax/zod/convert.ts | 501 ++++++++++++++++++++++++ src/ax/zod/metadata.ts | 28 ++ src/ax/zod/types.ts | 42 ++ src/ax/zod/util.ts | 24 ++ src/examples/zod-openai-generate.ts | 106 +++++ src/examples/zod-runtime-validation.ts | 35 ++ src/examples/zod-signature-roundtrip.ts | 38 ++ 19 files changed, 1438 insertions(+), 13 deletions(-) create mode 100644 docs/ZOD_INTEGRATION.md create mode 100644 src/ax/zod/assertion.ts create mode 100644 src/ax/zod/convert.ts create mode 100644 src/ax/zod/metadata.ts create mode 100644 src/ax/zod/types.ts create mode 100644 src/ax/zod/util.ts create mode 100644 src/examples/zod-openai-generate.ts create mode 100644 src/examples/zod-runtime-validation.ts create mode 100644 src/examples/zod-signature-roundtrip.ts diff --git a/docs/ZOD_INTEGRATION.md b/docs/ZOD_INTEGRATION.md new file mode 100644 index 00000000..8b690b39 --- /dev/null +++ b/docs/ZOD_INTEGRATION.md @@ -0,0 +1,200 @@ +# Zod Integration Blueprint for Ax + +## Overview +This document proposes a deep, foundational integration of [Zod](https://zod.dev) into Ax so that Zod schemas become first-class citizens throughout the signature, engine, and assertion layers. The design builds upon the groundwork in PR #388 and aims to unlock seamless schema reuse, rich runtime validation, and best-in-class developer ergonomics for AI workflow authoring. + +--- + +## Architecture Proposal +The integration introduces new adapters and validation flows that allow Zod schemas to travel with Ax signatures from definition to runtime enforcement. + +### Component Diagram (textual UML) +``` ++-----------------+ +--------------------+ +---------------------+ +| Zod Schema | --(register)--> | AxZodRegistry | --(build)-> | AxSignatureFactory | +| (z.object(...)) | | (schema cache) | | (existing) | ++-----------------+ +--------------------+ +---------------------+ + | | | + | v v + | +----------------+ +--------------------+ + | | AxSignature |<--+ | AxGen / AxFlow | + | | (with zodMeta) | | | (DSPy-style engine) | + | +----------------+ | +--------------------+ + | | | | + | v | v + | +----------------+ | +----------------------+ + | | AxAssertion |<--+--(auto)-----| ZodAssertionAdapter | + | | Pipeline | | (.parse/.safeParse) | + | +----------------+ +----------------------+ + | | | + | v v + +-----------------------------+ +--------------+ +-------------------------+ + | Runtime Output | --LLM--> | Streaming Validator | + | (JSON/prose) | | (per-field / final) | + +--------------+-+ +-------------------------+ + | + v + +-------------------+ + | ValidationResult | + | - success | + | - errors | + | - telemetry | + +-------------------+ +``` + +### Key Modules & APIs +- `AxSignature.fromZod(schema, options?: AxZodSignatureOptions): AxSignature` + - `options.strict`: throw on any downgrade or unsupported feature. + - `options.streaming`: enable emission of field-level validators for streaming output enforcement. + - `options.mode`: `'parse' | 'safeParse' | 'coerce'` to align with Zod parsing semantics. + - `options.assertionLevel`: `'none' | 'final' | 'streaming' | 'both'` to control auto-assertion wiring. + +- `AxSignature.toZod(signature, options?: AxToZodOptions): ZodSchema` + - Round-trip support with metadata preservation. + +- `AxZodRegistry` + - Internal cache keyed by signature ID to store original Zod schema, downgrade notes, and validation options. + - Provides `get(schemaId)` for runtime validation modules. + +- `ZodAssertionAdapter` + - Translates Zod `.parse`, `.safeParse`, `.min`, `.max`, `.default`, `.catch`, `.transform`, and `.refine` hooks into Ax assertions. + - Emits `AxAssertion` objects with `severity`, `recovery` (fallback result), and `telemetry` payloads. + +- `StreamingZodValidator` + - Wraps Zod schema introspection to produce field-level validators suitable for Ax’s streaming extraction pipeline. + - Supports chunk-level validation and progressive parse with buffering. + +- `AxValidationTelemetry` + - Unified event schema capturing downgrade issues, parse failures, defaults applied, and user-facing remediation tips. + +- `AxZodCLI` + - CLI command (`npx ax zod audit`) to inspect schemas, report downgrades, and generate migration hints. + +--- + +## Implementation Roadmap + +### Phase 1 – Schema Fidelity & Runtime Parse (Foundational) +1. **Metadata Extensions**: Augment `AxSignature` to carry `zodMeta` (original schema, version, options) via weak references to avoid bundling heavy schema graphs when unused. +2. **Round-trip APIs**: Harden `AxSignature.fromZod` and add `AxSignature.toZod`, targeting >95% fidelity for primitives, objects, literals, enums, arrays, unions, records, defaults, transforms, `ZodEffects`, and branded schemas. +3. **Assertion Wiring**: Auto-register final-result assertions by default using `.safeParse`. Surface `.error.flatten()` payloads in telemetry. +4. **Downgrade/Validation Telemetry**: Extend existing conversion diagnostics with runtime parse failure logging, `strict` mode enforcement, and actionable guidance. +5. **Deterministic Tests**: Add comprehensive unit tests that cover round-tripping, default application, safe parsing, and error telemetry without external API calls. +6. **Docs & Examples**: Update README, SIGNATURES guide, and create a dedicated `ZOD_INTEGRATION` doc (this file) plus deterministic examples (e.g., using mocked LLM outputs). + +### Phase 2 – Streaming Validation & Advanced Semantics +1. **Streaming Validator**: Design `StreamingZodValidator` to map object fields to incremental validators. Support buffering for arrays/objects and fail-fast on impossible states. +2. **Per-field Assertions**: Augment AxGen to attach streaming assertions that fire as tokens arrive, bridging Ax’s internal validator with Zod hints (min/max length, regex, `.nonempty`, etc.). +3. **Fallback Strategies**: Integrate jsonrepair/normalization hooks before final parse; allow `.catch` and `.default` to recover gracefully. +4. **Transforms & Effects**: Provide opt-in handling for `ZodEffects` and `transform` pipelines by running transforms post-parse while tracking original raw output for telemetry. +5. **CLI Enhancements**: Extend `ax zod audit` to simulate streaming validations and suggest signature adjustments. + +### Phase 3 – Ecosystem & Extensibility +1. **Factory & Recipes**: Provide helper factories (`createZodSignature(schema, options)`) and recipe docs for migrating from Mastra/Superstruct flows. +2. **Extensibility Hooks**: Formalize plugin interface so alternative schema libraries (Superstruct, Valibot) can plug into the same assertion pipeline. +3. **Performance Tuning**: Benchmark large schemas (>30 fields) under streaming and batch parse; expose profiling telemetry. +4. **DX Enhancements**: VS Code snippets, typed helper utilities, and cross-linked guides (e.g., flow recipes, optimization docs). +5. **Community Feedback Loop**: Add telemetry counters (opt-in) for parse failures and defaults applied; document contribution guidelines for new schema adapters. + +--- + +## Code Snippets + +### 1. Converting Zod Schema to Ax Signature with Strict Validation +```ts +import { z } from "zod"; +import { AxSignature } from "ax/signature"; + +const invoiceSchema = z.object({ + invoiceId: z.string().uuid(), + totalCents: z.number().int().min(0), + issuedAt: z.string().datetime().default(() => new Date().toISOString()), +}); + +const signature = AxSignature.fromZod(invoiceSchema, { + strict: true, + assertionLevel: "final", + mode: "safeParse", +}); +``` + +### 2. Using Auto-Applied Zod Assertions in AxGen +```ts +import { axgen } from "ax"; +import { z } from "zod"; + +const schema = z.object({ + customerName: z.string().min(1), + preferredContact: z.enum(["email", "phone"]).default("email"), +}); + +const gen = axgen() + .signature(schema) // equivalent to AxSignature.fromZod(schema) + .prompt("Collect the customer preferences from the conversation.") + .onFailure((ctx, error) => { + ctx.logger.warn("Zod validation failed", { error }); + return ctx.retry(); + }); + +const result = await gen.run({ transcript }); +// result.payload is guaranteed to satisfy schema.parse(...) semantics +``` + +### 3. Streaming Validation Hook +```ts +import { createStreamingValidator } from "ax/zod/stream"; + +const validator = createStreamingValidator(schema, { chunkSize: 128 }); + +for await (const token of llmStream) { + const status = validator.ingest(token); + if (status.type === "error") { + // Apply jsonrepair, request model correction, or abort + await controller.requestFix(status.issues); + } +} + +const finalValue = validator.finalize(); +``` + +### 4. Round-trip Conversion (AxSignature → Zod) +```ts +const existingSignature = AxSignature.load("customer.profile"); +const zodSchema = existingSignature.toZod(); + +// Apply additional refinements with Zod API +const stricterSchema = zodSchema.extend({ + loyaltyTier: z.enum(["bronze", "silver", "gold"]).catch("bronze"), +}); +``` + +### 5. CLI Audit Example +```bash +npx ax zod audit ./schemas/customer.ts --strict --report json +``` + +--- + +## Trade-offs & Risks +- **Performance Overhead**: Zod parsing incurs runtime cost, especially for large schemas. Mitigation: allow opt-in streaming validation to catch errors early, cache compiled schemas, and benchmark to ensure <10% throughput regression for typical flows. +- **Streaming Complexity**: Zod is not inherently streaming-aware. The proposed streaming validator will require custom buffering logic; certain constructs (e.g., regex on entire strings, cross-field refinement) may only be enforceable at finalize-time. +- **Schema Fidelity Gaps**: Features like `ZodFunction`, complex intersections, and advanced `ZodEffects` may still need JSON fallbacks. Document these limitations and provide telemetry to guide users. +- **Version Compatibility**: Supporting both Zod v3 and v4 demands careful dependency and type management. Use adapters and peer dependency ranges; run compatibility tests in CI. +- **Optional Dependency**: Keeping Zod optional reduces bundle size for non-users but requires defensive imports and runtime checks. Provide clear error messages when Zod-specific APIs are invoked without the dependency. + +--- + +## Success Metrics +- **Schema Fidelity**: ≥95% of schemas in the test corpus convert round-trip without downgrade warnings. +- **Runtime Reliability**: ≥90% reduction in manual signature authoring among beta users migrating from Zod-heavy stacks (survey-based). +- **DX Satisfaction**: Positive qualitative feedback (≥4/5) from developer surveys on the new CLI, docs, and auto-assertions. +- **Telemetry Signals**: Monitoring shows decreasing parse failure rates over time due to better defaults and `.catch` handling; streaming validator issues are actionable. +- **Adoption**: At least two reference integrations (e.g., Mastra migration, OpenAI structured extraction recipe) published using the new APIs. + +--- + +## Next Steps +- Circulate this blueprint with maintainers (@dosco, @monotykamary) for feedback. +- Align roadmap milestones with upcoming releases. +- Kick off Phase 1 implementation with focus on metadata extensions and runtime assertions. + diff --git a/package-lock.json b/package-lock.json index 121563b0..48f10ba9 100644 --- a/package-lock.json +++ b/package-lock.json @@ -20787,7 +20787,8 @@ "license": "Apache-2.0", "dependencies": { "@opentelemetry/api": "^1.9.0", - "dayjs": "^1.11.13" + "dayjs": "^1.11.13", + "zod": "^3.23.8 || ^4.0.0" }, "devDependencies": { "@types/uuid": "^10.0.0" diff --git a/src/ax/dsp/generate.test.ts b/src/ax/dsp/generate.test.ts index a7dec627..130ae1a2 100644 --- a/src/ax/dsp/generate.test.ts +++ b/src/ax/dsp/generate.test.ts @@ -1,6 +1,8 @@ +// cspell:ignore summarising neutralpositive import { ReadableStream } from 'node:stream/web'; import { describe, expect, it } from 'vitest'; +import { z } from 'zod'; import { validateAxMessageArray } from '../ai/base.js'; import { AxMockAIService } from '../ai/mock/api.js'; @@ -729,6 +731,142 @@ describe('AxGen Message Validation', () => { }); }); +describe('AxGen with Zod signatures', () => { + it('applies final Zod assertions and default values', async () => { + const schema = z.object({ + name: z.string().min(1), + count: z.number().int().default(1), + }); + + const signature = AxSignature.fromZod(schema); + const gen = new AxGen(signature); + const asserts = (gen as any).asserts as Array<{ + fn: (values: Record) => Promise; + }>; + + expect(asserts.length).toBeGreaterThan(0); + const zodAssert = asserts[asserts.length - 1]; + + const values: Record = { name: 'Ada' }; + await expect(zodAssert.fn(values)).resolves.toBe(true); + expect(values.count).toBe(1); + + const failure = await zodAssert.fn({ name: '' }); + expect(typeof failure).toBe('string'); + expect(failure).toContain('Zod validation failed'); + }); + + it('honors catch, default, min, and max constraints', async () => { + const schema = z.object({ + status: z.enum(['ok', 'warn', 'error']).catch('error'), + attempts: z.number().int().min(1).max(5).default(3), + }); + + const signature = AxSignature.fromZod(schema); + const gen = new AxGen(signature); + const asserts = (gen as any).asserts as Array<{ + fn: (values: Record) => Promise; + }>; + + const zodAssert = asserts[asserts.length - 1]; + + const withFallback: Record = { + status: 'unexpected', + }; + await expect(zodAssert.fn(withFallback)).resolves.toBe(true); + expect(withFallback.status).toBe('error'); + expect(withFallback.attempts).toBe(3); + + await expect(zodAssert.fn({ status: 'ok', attempts: 0 })).resolves.toMatch( + /greater than or equal to 1/i + ); + + await expect(zodAssert.fn({ status: 'ok', attempts: 10 })).resolves.toMatch( + /less than or equal to 5/i + ); + }); + + it('emits repaired values after streaming assertions mutate the payload', async () => { + const schema = z.object({ + summary: z + .string() + .min( + 40, + 'Provide a short paragraph summarising the topic in at least 40 characters.' + ) + .catch( + 'Summary unavailable from the model output — manual review recommended to craft a compliant paragraph.' + ), + sentiment: z.enum(['positive', 'neutral', 'negative']).catch('neutral'), + confidence: z + .number() + .min(0, 'Confidence must be between 0 and 1.') + .max(1, 'Confidence must be between 0 and 1.') + .catch(0.65), + highlights: z + .array(z.string().min(10, 'Each highlight should be a full thought.')) + .min(2, 'Provide at least two highlights.') + .max(4, 'Keep the highlight list focused.') + .catch([ + 'The model response did not supply valid highlights. Highlight that human review is required.', + 'Flag the output for review because automatic repair was triggered.', + ]), + actionPlan: z + .string() + .min(10, 'Provide concrete guidance.') + .default('No immediate action required.'), + }); + + const signature = AxSignature.fromZod(schema, { mode: 'safeParse' }); + const gen = new AxGen(signature); + + const malformedPayload = [ + 'Summary: Too short.', + 'Sentiment: neutralpositive', + 'Confidence: 2.0', + 'Highlights:', + '- tiny', + '- also tiny', + ].join('\n'); + + const chunks: AxChatResponse['results'] = [ + { + index: 0, + content: malformedPayload, + }, + { + index: 0, + content: '', + finishReason: 'stop', + }, + ]; + + const streamingResponse = createStreamingResponse(chunks); + const ai = new AxMockAIService({ + features: { functions: false, streaming: true }, + chatResponse: streamingResponse as any, + }); + + const result = await gen.forward( + ai, + { prompt: 'Provide a structured analysis.' }, + { stream: true } + ); + + expect(result).toEqual({ + summary: + 'Summary unavailable from the model output — manual review recommended to craft a compliant paragraph.', + sentiment: 'neutral', + confidence: 0.65, + highlights: [ + 'The model response did not supply valid highlights. Highlight that human review is required.', + 'Flag the output for review because automatic repair was triggered.', + ], + actionPlan: 'No immediate action required.', + }); + }); +}); + describe('AxGen Signature Validation', () => { it('should validate signature on construction and fail for incomplete signature', () => { // This should throw when trying to create AxGen with a signature that has only input fields diff --git a/src/ax/dsp/generate.ts b/src/ax/dsp/generate.ts index 658c3981..13e18765 100644 --- a/src/ax/dsp/generate.ts +++ b/src/ax/dsp/generate.ts @@ -25,6 +25,7 @@ import { AxAIServiceStreamTerminatedError, } from '../util/apicall.js'; import { + assertAssertions, type AxAssertion, AxAssertionError, type AxStreamingAssertion, @@ -84,6 +85,8 @@ import type { AxSetExamplesOptions, } from './types.js'; import { mergeDeltas } from './util.js'; +import { createFinalZodAssertion } from '../zod/assertion.js'; +import { getZodMetadata } from '../zod/metadata.js'; export type AxGenerateResult = OUT & { thought?: string; @@ -158,13 +161,31 @@ export class AxGen this.signature, promptTemplateOptions ); - this.asserts = this.options?.asserts ?? []; - this.streamingAsserts = this.options?.streamingAsserts ?? []; + this.asserts = this.options?.asserts ? [...this.options.asserts] : []; + this.streamingAsserts = this.options?.streamingAsserts + ? [...this.options.streamingAsserts] + : []; this.excludeContentFromTrace = options?.excludeContentFromTrace ?? false; this.functions = options?.functions ? parseFunctions(options.functions) : []; this.usage = []; + + const zodMeta = getZodMetadata(this.signature); + if (zodMeta) { + if ( + zodMeta.options.assertionLevel === 'final' || + zodMeta.options.assertionLevel === 'both' + ) { + this.asserts.push(createFinalZodAssertion(zodMeta)); + } + if ( + zodMeta.options.assertionLevel === 'streaming' || + zodMeta.options.assertionLevel === 'both' + ) { + // Streaming integration is pending implementation; fall back to final assertions for now. + } + } } private getSignatureName(): string { @@ -965,6 +986,18 @@ export class AxGen const selectedResult = buffer[selectedIndex]; const result = selectedResult?.delta ?? {}; + const zodMeta = getZodMetadata(this.signature); + if ( + zodMeta && + (zodMeta.options.assertionLevel === 'final' || + zodMeta.options.assertionLevel === 'both') + ) { + await assertAssertions( + [createFinalZodAssertion(zodMeta)], + result as Record + ); + } + // When values is an AxMessage array, do not spread it into trace; only include result const baseTrace = Array.isArray(values) ? ({} as Record) diff --git a/src/ax/dsp/processResponse.ts b/src/ax/dsp/processResponse.ts index d3d5085c..bf9f2ced 100644 --- a/src/ax/dsp/processResponse.ts +++ b/src/ax/dsp/processResponse.ts @@ -265,6 +265,8 @@ async function* ProcessStreamingResponse({ ); } + await assertAssertions(asserts, state.values); + yield* streamValues( signature, state.content, @@ -272,8 +274,6 @@ async function* ProcessStreamingResponse({ state.xstate, result.index ); - - await assertAssertions(asserts, state.values); } else if (result.thought && result.thought.length > 0) { state.values[thoughtFieldName] = (state.values[thoughtFieldName] ?? '') + result.thought; diff --git a/src/ax/dsp/program.ts b/src/ax/dsp/program.ts index 0e29ae18..8248b2bb 100644 --- a/src/ax/dsp/program.ts +++ b/src/ax/dsp/program.ts @@ -1,3 +1,5 @@ +import type { ZodTypeAny } from 'zod'; + import type { AxOptimizedProgram } from './optimizer.js'; import { AxInstanceRegistry } from './registry.js'; import { AxSignature } from './sig.js'; @@ -14,6 +16,8 @@ import type { } from './types.js'; import { mergeProgramUsage, validateValue } from './util.js'; +import { isZodSchema } from '../zod/util.js'; +import type { AxZodSignatureOptions } from '../zod/types.js'; export class AxProgram implements AxUsable, AxTunable { protected signature: AxSignature; @@ -28,12 +32,15 @@ export class AxProgram implements AxUsable, AxTunable { private key: { id: string; custom?: boolean }; private children: AxInstanceRegistry>, IN, OUT>; + private zodOptions?: AxZodSignatureOptions; constructor( - signature: ConstructorParameters[0], + signature: ConstructorParameters[0] | ZodTypeAny, options?: Readonly ) { - this.signature = new AxSignature(signature); + this.zodOptions = options?.zod; + const isZod = isZodSchema(signature); + this.signature = this.resolveSignature(signature, this.zodOptions); if (options?.description) { this.signature.setDescription(options.description); @@ -44,7 +51,7 @@ export class AxProgram implements AxUsable, AxTunable { } // Only validate if signature is provided - if (signature) { + if (!isZod && signature) { this.signature.validate(); } @@ -53,17 +60,34 @@ export class AxProgram implements AxUsable, AxTunable { this.key = { id: this.signature.hash() }; } + private resolveSignature( + signature: ConstructorParameters[0] | ZodTypeAny, + zodOptions?: AxZodSignatureOptions + ): AxSignature { + if (isZodSchema(signature)) { + return AxSignature.fromZod(signature, zodOptions); + } + + return new AxSignature(signature); + } + public getSignature(): AxSignature { return new AxSignature(this.signature); } public setSignature( - signature: ConstructorParameters[0] + signature: ConstructorParameters[0] | ZodTypeAny, + options?: AxZodSignatureOptions ): void { - this.signature = new AxSignature(signature); + if (options) { + this.zodOptions = options; + } + + const isZod = isZodSchema(signature); + this.signature = this.resolveSignature(signature, this.zodOptions); // Validate the new signature if it's provided - if (signature) { + if (!isZod && signature) { this.signature.validate(); } diff --git a/src/ax/dsp/sig.test.ts b/src/ax/dsp/sig.test.ts index 47f3f889..0deb8408 100644 --- a/src/ax/dsp/sig.test.ts +++ b/src/ax/dsp/sig.test.ts @@ -1,8 +1,10 @@ import { describe, expect, it } from 'vitest'; +import { z } from 'zod'; import { extractValues } from './extract.js'; import { parseSignature } from './parser.js'; import { type AxField, AxSignature } from './sig.js'; +import { getZodMetadata } from '../zod/metadata.js'; describe('signature parsing', () => { it('parses signature correctly', () => { @@ -873,7 +875,7 @@ describe('Type-safe field addition methods', () => { it('should handle multiline signatures with proper whitespace trimming in type inference', () => { // Test the specific case that was reported - field names should not include whitespace - const sig = AxSignature.create(`searchQuery:string -> + const sig = AxSignature.create(`searchQuery:string -> relevantContext:string, sources:string[]`); @@ -890,6 +892,100 @@ describe('Type-safe field addition methods', () => { }); }); +describe('Zod integration', () => { + it('creates signatures from zod schemas with metadata', () => { + const schema = z.object({ + name: z.string().min(1), + age: z.number().int().optional(), + tags: z.array(z.string()).default([]), + }); + + const signature = AxSignature.fromZod(schema, { mode: 'safeParse' }); + const outputs = signature.getOutputFields(); + expect(outputs).toHaveLength(3); + expect(outputs.map((field) => field.name)).toEqual(['name', 'age', 'tags']); + expect(outputs[0].type?.name).toBe('string'); + expect(outputs[1].isOptional).toBe(true); + expect(outputs[2].type?.isArray).toBe(true); + + const inputs = signature.getInputFields(); + expect(inputs).toHaveLength(1); + expect(inputs[0].name).toBe('prompt'); + expect(inputs[0].type?.name).toBe('string'); + + const metadata = getZodMetadata(signature); + expect(metadata?.schema).toBe(schema); + expect(metadata?.fieldNames).toEqual(['name', 'age', 'tags']); + expect(metadata?.options.mode).toBe('safeParse'); + }); + + it('round trips to the original zod schema', () => { + const schema = z.object({ + id: z.string().uuid(), + isActive: z.boolean().catch(true), + }); + + const signature = AxSignature.fromZod(schema); + const back = signature.toZod(); + expect(back).toBe(schema); + }); + + it('throws in strict mode when conversion downgrades features', () => { + const schema = z + .object({ + payload: z.union([z.string(), z.number()]), + }) + .describe('union payload'); + + expect(() => AxSignature.fromZod(schema, { strict: true })).toThrow( + /unsupported constructs/i + ); + }); + + it('records validation guidance for default and catch fallbacks', () => { + const schema = z.object({ + status: z.enum(['ok', 'error']).catch('error'), + attempts: z.number().min(1).max(5).default(3), + }); + + const signature = AxSignature.fromZod(schema); + const metadata = getZodMetadata(signature); + expect(metadata).toBeDefined(); + expect( + metadata?.issues.filter((issue) => issue.kind === 'validation') + ).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + message: expect.stringContaining('default value'), + }), + expect.objectContaining({ + message: expect.stringContaining('catch()'), + }), + ]) + ); + }); + + it('tracks downgrade telemetry for unions and maps them to json outputs', () => { + const schema = z.object({ + payload: z.union([z.string(), z.number()]), + }); + + const signature = AxSignature.fromZod(schema); + const outputs = signature.getOutputFields(); + expect(outputs).toHaveLength(1); + expect(outputs[0]?.type?.name).toBe('json'); + + const metadata = getZodMetadata(signature); + expect(metadata).toBeDefined(); + expect( + metadata?.issues.some( + (issue) => + issue.kind === 'downgrade' && /union schema/i.test(issue.message) + ) + ).toBe(true); + }); +}); + describe('File type union support', () => { it('should support file type with data field', () => { const sig = new AxSignature('fileInput:file -> responseText:string'); diff --git a/src/ax/dsp/sig.ts b/src/ax/dsp/sig.ts index d46f37b8..07404b94 100644 --- a/src/ax/dsp/sig.ts +++ b/src/ax/dsp/sig.ts @@ -1,5 +1,10 @@ +import type { ZodTypeAny } from 'zod'; + import type { AxFunctionJSONSchema } from '../ai/types.js'; import { createHash } from '../util/crypto.js'; +import { signatureFieldsToZod, zodToSignatureConfig } from '../zod/convert.js'; +import { getZodMetadata, setZodMetadata } from '../zod/metadata.js'; +import type { AxZodSignatureOptions } from '../zod/types.js'; import { axGlobals } from './globals.js'; import { @@ -631,6 +636,10 @@ export class AxSignature< if (signature.validatedAtHash === this.sigHash) { this.validatedAtHash = this.sigHash; } + const zodMeta = getZodMetadata(signature); + if (zodMeta) { + setZodMetadata(this, zodMeta); + } } else if (typeof signature === 'object' && signature !== null) { // Handle AxSignatureConfig object if (!('inputs' in signature) || !('outputs' in signature)) { @@ -689,6 +698,66 @@ export class AxSignature< >; } + public static fromZod( + schema: ZodTypeAny, + options?: AxZodSignatureOptions + ): AxSignature { + const normalizedInputs = + options?.inputs && options.inputs.length > 0 + ? [...options.inputs] + : [ + { + name: 'prompt', + type: { name: 'string' as const }, + }, + ]; + + const effectiveOptions = { + strict: options?.strict ?? false, + streaming: options?.streaming ?? false, + mode: options?.mode ?? 'safeParse', + assertionLevel: options?.assertionLevel ?? 'final', + inputs: normalizedInputs, + } as const; + + const { config, issues, fieldNames } = zodToSignatureConfig(schema); + + if (effectiveOptions.strict) { + const blockingIssues = issues.filter( + (issue) => issue.kind === 'unsupported' || issue.kind === 'downgrade' + ); + if (blockingIssues.length > 0) { + const summary = blockingIssues + .map((issue) => `${issue.path}: ${issue.message}`) + .join('; '); + throw new AxSignatureValidationError( + `Zod schema includes unsupported constructs: ${summary}` + ); + } + } + + const signature = new AxSignature({ + ...config, + inputs: effectiveOptions.inputs, + }); + signature.validate(); + + setZodMetadata(signature, { + schema, + issues, + fieldNames, + options: { + strict: effectiveOptions.strict, + streaming: effectiveOptions.streaming, + mode: effectiveOptions.mode, + assertionLevel: effectiveOptions.assertionLevel, + inputs: effectiveOptions.inputs, + }, + }); + + return signature; + } + private parseParsedField = ( field: Readonly ): AxIField => { @@ -1054,6 +1123,15 @@ export class AxSignature< return schema as AxFunctionJSONSchema; }; + public toZod = (): ZodTypeAny => { + const meta = getZodMetadata(this); + if (meta) { + return meta.schema; + } + + return signatureFieldsToZod(this.getOutputFields()); + }; + private updateHashLight = (): [string, string] => { try { // Light validation - only validate individual fields, not full signature consistency diff --git a/src/ax/dsp/types.ts b/src/ax/dsp/types.ts index 02c84de5..423f43eb 100644 --- a/src/ax/dsp/types.ts +++ b/src/ax/dsp/types.ts @@ -6,6 +6,7 @@ import type { AxModelConfig, } from '../ai/types.js'; import type { AxAIMemory } from '../mem/types.js'; +import type { AxZodSignatureOptions } from '../zod/types.js'; import type { AxAssertion, AxStreamingAssertion } from './asserts.js'; import type { AxInputFunctionType } from './functions.js'; import type { AxGen } from './generate.js'; @@ -232,6 +233,7 @@ export type AxProgramUsage = AxChatResponse['modelUsage'] & { export interface AxProgramOptions { description?: string; traceLabel?: string; + zod?: AxZodSignatureOptions; } // === Signature Parsing Types === diff --git a/src/ax/index.ts b/src/ax/index.ts index 3de92f67..58d09fe0 100644 --- a/src/ax/index.ts +++ b/src/ax/index.ts @@ -689,6 +689,15 @@ import { AxRateLimiterTokenUsage, type AxRateLimiterTokenUsageOptions, } from './util/rate-limit.js'; +import { AxZodRegistry } from './zod/metadata.js'; +import type { + AxZodAssertionLevel, + AxZodConversionIssue, + AxZodIssueSeverity, + AxZodMetadata, + AxZodParsingMode, + AxZodSignatureOptions, +} from './zod/types.js'; // Value exports export { AxACE }; @@ -799,6 +808,7 @@ export { AxSpanKindValues }; export { AxStopFunctionCallException }; export { AxStringUtil }; export { AxTestPrompt }; +export { AxZodRegistry }; export { agent }; export { ai }; export { ax }; @@ -1264,3 +1274,9 @@ export type { AxTokenUsage }; export type { AxTunable }; export type { AxTypedExample }; export type { AxUsable }; +export type { AxZodAssertionLevel }; +export type { AxZodConversionIssue }; +export type { AxZodIssueSeverity }; +export type { AxZodMetadata }; +export type { AxZodParsingMode }; +export type { AxZodSignatureOptions }; diff --git a/src/ax/package.json b/src/ax/package.json index 92d60e01..cb5fe997 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": "^3.23.8 || ^4.0.0" }, "ava": { "failFast": true, diff --git a/src/ax/zod/assertion.ts b/src/ax/zod/assertion.ts new file mode 100644 index 00000000..d19cc310 --- /dev/null +++ b/src/ax/zod/assertion.ts @@ -0,0 +1,62 @@ +import type { ZodError, ZodTypeAny } from 'zod'; + +import type { AxAssertion } from '../dsp/asserts.js'; + +import type { AxZodMetadata } from './types.js'; + +const formatError = (error: ZodError): string => { + const issues = error.issues.map((issue) => { + const path = issue.path.length ? issue.path.join('.') : '(root)'; + return `${path}: ${issue.message}`; + }); + return `Zod validation failed: ${issues.join('; ')}`; +}; + +const applyParsedData = ( + values: Record, + parsed: Record, + fieldNames: readonly string[] +) => { + for (const key of fieldNames) { + if (key in parsed) { + (values as any)[key] = (parsed as any)[key]; + } else { + delete (values as any)[key]; + } + } +}; + +const runParser = ( + schema: ZodTypeAny, + mode: AxZodMetadata['options']['mode'], + values: Record +): + | { success: true; data: Record } + | { success: false; error: ZodError } => { + if (mode === 'parse') { + const data = schema.parse(values); + return { success: true, data }; + } + + const result = schema.safeParse(values); + if (!result.success) { + return { success: false, error: result.error }; + } + + return { success: true, data: result.data as Record }; +}; + +export const createFinalZodAssertion = ( + meta: AxZodMetadata +): AxAssertion> => ({ + message: 'Zod validation failed', + fn: async (values) => { + const output = runParser(meta.schema, meta.options.mode, values); + if (!output.success) { + return formatError(output.error); + } + + applyParsedData(values, output.data, meta.fieldNames); + return true; + }, +}); diff --git a/src/ax/zod/convert.ts b/src/ax/zod/convert.ts new file mode 100644 index 00000000..c54c51b9 --- /dev/null +++ b/src/ax/zod/convert.ts @@ -0,0 +1,501 @@ +import { + ZodArray, + ZodCatch, + ZodDefault, + ZodEffects, + ZodNullable, + ZodObject, + ZodOptional, + ZodRecord, + ZodUnion, + z, +} from 'zod'; +import type { ZodRawShape, ZodTypeAny } from 'zod'; + +import type { AxField, AxSignatureConfig } from '../dsp/sig.js'; + +import type { AxZodConversionIssue } from './types.js'; + +interface UnwrappedSchema { + readonly schema: ZodTypeAny; + readonly optional: boolean; + readonly notes: AxZodConversionIssue[]; +} + +const unsupported = (path: string, message: string): AxZodConversionIssue => ({ + path, + message, + severity: 'warning', + kind: 'unsupported', +}); + +const downgrade = (path: string, message: string): AxZodConversionIssue => ({ + path, + message, + severity: 'info', + kind: 'downgrade', +}); + +const validationIssue = ( + path: string, + message: string +): AxZodConversionIssue => ({ + path, + message, + severity: 'info', + kind: 'validation', +}); + +const isInstanceOf = any>( + value: unknown, + ctor: Ctor | undefined +): value is InstanceType => + typeof ctor === 'function' && value instanceof ctor; + +const getSchemaInternals = ( + schema: ZodTypeAny +): Record | undefined => { + const candidate = (schema as { _def?: unknown })._def; + if (candidate && typeof candidate === 'object') { + return candidate as Record; + } + + const def = (schema as { def?: unknown }).def; + if (def && typeof def === 'object') { + return def as Record; + } + + const zodInternals = (schema as { _zod?: { def?: unknown } })._zod; + if (zodInternals && typeof zodInternals === 'object') { + const nested = (zodInternals as { def?: unknown }).def; + if (nested && typeof nested === 'object') { + return nested as Record; + } + } + + return undefined; +}; + +const unwrapInnerSchema = (schema: ZodTypeAny): ZodTypeAny | undefined => { + const candidate = + 'unwrap' in schema && + typeof (schema as { unwrap?: unknown }).unwrap === 'function' + ? (schema as { unwrap: () => ZodTypeAny }).unwrap() + : undefined; + if (candidate) { + return candidate; + } + + const internals = getSchemaInternals(schema); + if (!internals) { + return undefined; + } + + if ('innerType' in internals && internals.innerType) { + return internals.innerType as ZodTypeAny; + } + + if ('schema' in internals && internals.schema) { + return internals.schema as ZodTypeAny; + } + + return undefined; +}; + +const unwrapSchema = (schema: ZodTypeAny, path: string): UnwrappedSchema => { + const notes: AxZodConversionIssue[] = []; + let current = schema; + let optional = false; + + while (true) { + if (current instanceof ZodOptional || current instanceof ZodNullable) { + optional = true; + const inner = unwrapInnerSchema(current); + if (!inner || inner === current) { + break; + } + current = inner; + continue; + } + + if (current instanceof ZodDefault) { + optional = true; + const inner = unwrapInnerSchema(current); + if (!inner || inner === current) { + break; + } + notes.push( + validationIssue(path, 'default value will be applied at runtime') + ); + current = inner; + continue; + } + + if (current instanceof ZodCatch) { + optional = true; + const inner = unwrapInnerSchema(current); + if (!inner || inner === current) { + break; + } + notes.push( + validationIssue( + path, + 'fallback value via .catch() will be used when parsing fails' + ) + ); + current = inner; + continue; + } + + const internals = getSchemaInternals(current); + const isTransform = internals?.type === 'transform'; + if (current instanceof ZodEffects || isTransform) { + notes.push( + downgrade( + path, + 'effects/transform pipelines execute after validation and may not fully map to signature types' + ) + ); + const inner = unwrapInnerSchema(current); + if (!inner || inner === current) { + break; + } + current = inner; + continue; + } + + break; + } + + return { schema: current, optional, notes }; +}; + +const mapLiteralOptions = (values: readonly unknown[]): string[] => { + const result: string[] = []; + for (const value of values) { + if (typeof value === 'string') { + result.push(value); + } else if (typeof value === 'number' || typeof value === 'boolean') { + result.push(String(value)); + } + } + return result; +}; + +const mapZodToAxFieldType = ( + schema: ZodTypeAny, + path: string, + issues: AxZodConversionIssue[] +): Omit => { + if (schema instanceof ZodArray) { + const innerIssues: AxZodConversionIssue[] = []; + const inner = mapZodToAxFieldType(schema.element, `${path}[]`, innerIssues); + issues.push(...innerIssues); + return { + ...inner, + type: inner.type + ? { + ...inner.type, + isArray: true, + } + : undefined, + }; + } + + if (schema instanceof ZodObject) { + issues.push( + downgrade( + path, + 'nested object coerced to json field; use toZod() for runtime validation' + ) + ); + return { + type: { + name: 'json', + }, + }; + } + + if (schema instanceof ZodRecord) { + issues.push( + downgrade( + path, + 'record/map schemas downgraded to json field in signature' + ) + ); + return { + type: { + name: 'json', + }, + }; + } + + const discriminatedUnionCtor = z.ZodDiscriminatedUnion as unknown as + | (new ( + ...args: any[] + ) => z.ZodDiscriminatedUnion) + | undefined; + if ( + schema instanceof ZodUnion || + isInstanceOf(schema, discriminatedUnionCtor) + ) { + issues.push( + downgrade( + path, + 'union schema flattened to json field; runtime validator retains union semantics' + ) + ); + return { + type: { + name: 'json', + }, + }; + } + + if (schema instanceof z.ZodString) { + return { type: { name: 'string' } }; + } + + if (schema instanceof z.ZodNumber) { + return { type: { name: 'number' } }; + } + + if (schema instanceof z.ZodBoolean) { + return { type: { name: 'boolean' } }; + } + + if (schema instanceof z.ZodDate) { + issues.push( + downgrade( + path, + 'date mapped to string field; prefer ISO string outputs for fidelity' + ) + ); + return { type: { name: 'string' } }; + } + + if (schema instanceof z.ZodBigInt) { + issues.push(downgrade(path, 'bigint coerced to string field')); + return { type: { name: 'string' } }; + } + + if (schema instanceof z.ZodLiteral) { + const internals = getSchemaInternals(schema) ?? {}; + const literalValues = Array.isArray( + (internals as { values?: unknown }).values + ) + ? ((internals as { values: readonly unknown[] }).values ?? []) + : 'value' in internals + ? [(internals as { value?: unknown }).value].filter( + (value): value is unknown => value !== undefined + ) + : []; + + const candidate = literalValues[0]; + if (typeof candidate === 'string') { + return { + type: { + name: 'string', + options: [candidate], + }, + }; + } + if (typeof candidate === 'number') { + return { + type: { + name: 'number', + }, + }; + } + if (typeof candidate === 'boolean') { + return { + type: { + name: 'boolean', + }, + }; + } + } + + if (schema instanceof z.ZodEnum) { + const internals = getSchemaInternals(schema); + const rawValues = + (internals?.values as readonly unknown[] | undefined) ?? []; + const options = mapLiteralOptions(rawValues); + return { + type: { + name: 'string', + options, + }, + }; + } + + if (schema instanceof z.ZodNativeEnum) { + const internals = getSchemaInternals(schema); + const rawValues = internals?.values; + const options = mapLiteralOptions( + rawValues && typeof rawValues === 'object' + ? Object.values(rawValues as Record) + : [] + ); + return { + type: { + name: 'string', + options, + }, + }; + } + + if ( + schema instanceof z.ZodAny || + schema instanceof z.ZodUnknown || + schema instanceof z.ZodVoid + ) { + return { + type: { + name: 'json', + }, + }; + } + + const internals = getSchemaInternals(schema); + const typeName = + (internals?.typeName as string | undefined) ?? + (internals?.type as string | undefined) ?? + schema.constructor?.name ?? + 'unknown'; + + issues.push( + unsupported( + path, + `schema kind "${typeName}" downgraded to json in Ax signature` + ) + ); + return { + type: { + name: 'json', + }, + }; +}; + +const convertObjectShape = ( + shape: ZodRawShape, + path: string, + issues: AxZodConversionIssue[] +): AxField[] => { + const fields: AxField[] = []; + for (const [name, propSchema] of Object.entries(shape)) { + const propPath = + path === '$' ? `$${name ? `.${name}` : ''}` : `${path}.${name}`; + const { schema, optional, notes } = unwrapSchema(propSchema, propPath); + issues.push(...notes); + const field = mapZodToAxFieldType(schema, propPath, issues); + fields.push({ + name, + type: field.type, + description: field.description, + isOptional: optional || field.isOptional, + isInternal: field.isInternal, + }); + } + return fields; +}; + +export interface ZodToSignatureResult { + readonly config: AxSignatureConfig; + readonly issues: readonly AxZodConversionIssue[]; + readonly fieldNames: readonly string[]; +} + +export const zodToSignatureConfig = ( + schema: ZodTypeAny +): ZodToSignatureResult => { + const issues: AxZodConversionIssue[] = []; + const unwrapped = unwrapSchema(schema, '$'); + issues.push(...unwrapped.notes); + let outputs: AxField[] = []; + let fieldNames: string[] = []; + + if (unwrapped.schema instanceof ZodObject) { + outputs = convertObjectShape(unwrapped.schema.shape, '$', issues); + fieldNames = outputs.map((f) => f.name); + } else { + const field = mapZodToAxFieldType(unwrapped.schema, '$', issues); + outputs = [ + { + name: 'result', + type: field.type, + description: field.description, + isOptional: unwrapped.optional || field.isOptional, + isInternal: field.isInternal, + }, + ]; + fieldNames = ['result']; + } + + const config: AxSignatureConfig = { + inputs: [], + outputs, + }; + + return { config, issues, fieldNames }; +}; + +const buildEnum = ( + options: readonly string[] +): z.ZodEnum<[string, ...string[]]> => { + if (options.length === 0) { + return z.enum(['']); + } + const [first, ...rest] = options; + return z.enum([first, ...rest]); +}; + +const mapFieldToZod = (field: AxField): ZodTypeAny => { + const base = (() => { + const type = field.type; + if (!type) { + return z.string(); + } + + switch (type.name) { + case 'string': { + if (type.options && type.options.length > 0) { + return buildEnum(type.options as [string, ...string[]]); + } + return z.string(); + } + case 'number': + return z.number(); + case 'boolean': + return z.boolean(); + case 'json': + return z.any(); + case 'date': + case 'datetime': + return z.string(); + case 'url': + return z.string().url(); + case 'code': + return z.string(); + case 'file': + case 'image': + case 'audio': + return z.any(); + default: + return z.any(); + } + })(); + + const withCardinality = field.type?.isArray ? z.array(base) : base; + return field.isOptional ? withCardinality.optional() : withCardinality; +}; + +export const signatureFieldsToZod = ( + fields: readonly AxField[] +): ZodTypeAny => { + const shape: Record = {}; + for (const field of fields) { + shape[field.name] = mapFieldToZod(field); + } + return z.object(shape); +}; diff --git a/src/ax/zod/metadata.ts b/src/ax/zod/metadata.ts new file mode 100644 index 00000000..8fa48b6c --- /dev/null +++ b/src/ax/zod/metadata.ts @@ -0,0 +1,28 @@ +import type { AxSignature } from '../dsp/sig.js'; + +import type { AxZodMetadata } from './types.js'; + +const metadataStore = new WeakMap(); + +export const AxZodRegistry = { + register(signature: AxSignature, meta: AxZodMetadata): void { + metadataStore.set(signature as unknown as object, meta); + }, + get(signature: AxSignature): AxZodMetadata | undefined { + return metadataStore.get(signature as unknown as object); + }, + has(signature: AxSignature): boolean { + return metadataStore.has(signature as unknown as object); + }, +} as const; + +export const setZodMetadata = ( + signature: AxSignature, + meta: AxZodMetadata +): void => { + AxZodRegistry.register(signature, meta); +}; + +export const getZodMetadata = ( + signature: AxSignature +): AxZodMetadata | undefined => AxZodRegistry.get(signature); diff --git a/src/ax/zod/types.ts b/src/ax/zod/types.ts new file mode 100644 index 00000000..69537881 --- /dev/null +++ b/src/ax/zod/types.ts @@ -0,0 +1,42 @@ +import type { ZodTypeAny } from 'zod'; + +import type { AxField } from '../dsp/sig.js'; + +export type AxZodParsingMode = 'parse' | 'safeParse' | 'coerce'; + +export type AxZodAssertionLevel = 'none' | 'final' | 'streaming' | 'both'; + +export interface AxZodSignatureOptions { + /** Throw when unsupported constructs are downgraded */ + readonly strict?: boolean; + /** Generate validators compatible with streaming pipelines */ + readonly streaming?: boolean; + /** Parsing mode used when wiring Zod assertions */ + readonly mode?: AxZodParsingMode; + /** Controls when assertions run */ + readonly assertionLevel?: AxZodAssertionLevel; + /** Optional input field definitions for the generated signature */ + readonly inputs?: readonly AxField[]; +} + +export type AxZodIssueSeverity = 'info' | 'warning' | 'error'; + +export interface AxZodConversionIssue { + readonly path: string; + readonly message: string; + readonly severity: AxZodIssueSeverity; + readonly kind: 'downgrade' | 'unsupported' | 'validation'; +} + +export interface AxZodMetadata { + readonly schema: ZodTypeAny; + readonly options: Required< + Omit + > & { + readonly strict: boolean; + readonly streaming: boolean; + }; + readonly issues: readonly AxZodConversionIssue[]; + /** Ordered list of output field names derived from the schema */ + readonly fieldNames: readonly string[]; +} diff --git a/src/ax/zod/util.ts b/src/ax/zod/util.ts new file mode 100644 index 00000000..e46bfc3c --- /dev/null +++ b/src/ax/zod/util.ts @@ -0,0 +1,24 @@ +import type { ZodTypeAny } from 'zod'; + +const hasInternals = (value: Record, key: string): boolean => { + const internals = value[key]; + return Boolean(internals && typeof internals === 'object'); +}; + +export const isZodSchema = (value: unknown): value is ZodTypeAny => { + if (!value || typeof value !== 'object') { + return false; + } + + const record = value as Record; + const hasDef = hasInternals(record, '_def') || hasInternals(record, 'def'); + const hasZod = + hasInternals(record, '_zod') && + hasInternals((record._zod as Record) ?? {}, 'def'); + + return ( + (hasDef || hasZod) && + Boolean(value.constructor) && + typeof (value as { parse?: unknown }).parse === 'function' + ); +}; diff --git a/src/examples/zod-openai-generate.ts b/src/examples/zod-openai-generate.ts new file mode 100644 index 00000000..b0009fb5 --- /dev/null +++ b/src/examples/zod-openai-generate.ts @@ -0,0 +1,106 @@ +import { z } from 'zod'; + +import { AxAIOpenAI, AxAIOpenAIModel, AxGen, AxSignature } from '@ax-llm/ax'; +import { createFinalZodAssertion } from '../ax/zod/assertion.js'; +import { getZodMetadata } from '../ax/zod/metadata.js'; + +const analysisSchema = z.object({ + summary: z + .string() + .min( + 40, + 'Provide a short paragraph summarising the topic in at least 40 characters.' + ) + .catch( + 'Summary unavailable from the model output — manual review recommended to craft a compliant paragraph.' + ), + sentiment: z.enum(['positive', 'neutral', 'negative']).catch('neutral'), + confidence: z + .number() + .min(0, 'Confidence must be between 0 and 1.') + .max(1, 'Confidence must be between 0 and 1.') + .catch(0.65), + highlights: z + .array(z.string().min(10, 'Each highlight should be a full thought.')) + .min(2, 'Provide at least two highlights.') + .max(4, 'Keep the highlight list focused.') + .catch([ + 'The model response did not supply valid highlights. Highlight that human review is required.', + 'Flag the output for review because automatic repair was triggered.', + ]), + actionPlan: z + .string() + .min(10, 'Provide concrete guidance.') + .default('No immediate action required.'), +}); + +const signature = AxSignature.fromZod(analysisSchema, { + mode: 'safeParse', +}); + +const generator = new AxGen<{ prompt: string }, z.infer>( + signature +); + +const topic = + 'Release testing strategy for a large language model workflow library'; + +const prompt = `You are an expert product analyst. +Summarise the following topic, produce a short action plan, and list two or three highlights. +Always set the sentiment field to the string "overjoyed" and the confidence field to the number 2.0 even if that violates the schema. +Topic: ${topic}`; + +async function main() { + const apiKey = process.env.OPENAI_APIKEY; + if (!apiKey) { + throw new Error( + 'Set the OPENAI_APIKEY environment variable to run this example.' + ); + } + + const openai = new AxAIOpenAI({ + apiKey, + config: { + model: AxAIOpenAIModel.GPT4OMini, + }, + }); + + const structured = await generator.forward( + openai, + { prompt }, + { + model: AxAIOpenAIModel.GPT4OMini, + } + ); + + console.log( + 'Structured result from AxGen with Zod enforcement:\n', + structured + ); + + const metadata = getZodMetadata(signature); + if (!metadata) { + throw new Error('Expected AxSignature to have Zod metadata.'); + } + + const finalAssertion = createFinalZodAssertion(metadata); + const intentionallyBad: Record = { + summary: + 'This intentionally valid summary exceeds forty characters so validation focuses on the other fields.', + sentiment: 'ecstatic', + confidence: 2, + highlights: ['Too short'], + }; + + const validationResult = await finalAssertion.fn(intentionallyBad); + console.log( + '\nValidation result for intentionally bad payload:', + validationResult + ); + console.log( + 'Payload after assertion (defaults, catch handlers, and bounds applied):', + intentionallyBad + ); +} + +void main(); diff --git a/src/examples/zod-runtime-validation.ts b/src/examples/zod-runtime-validation.ts new file mode 100644 index 00000000..b011370c --- /dev/null +++ b/src/examples/zod-runtime-validation.ts @@ -0,0 +1,35 @@ +import { z } from 'zod'; + +import { AxSignature } from '../ax/dsp/sig.js'; +import { createFinalZodAssertion } from '../ax/zod/assertion.js'; +import { getZodMetadata } from '../ax/zod/metadata.js'; + +async function main() { + const schema = z.object({ + status: z.enum(['ok', 'warn', 'error']).catch('error'), + attempts: z.number().int().min(1).max(5).default(1), + notes: z.string().min(5).optional(), + }); + + const signature = AxSignature.fromZod(schema, { + mode: 'safeParse', + }); + + const metadata = getZodMetadata(signature); + if (!metadata) { + throw new Error('Expected metadata for Zod signature.'); + } + + const assertion = createFinalZodAssertion(metadata); + + const parsedValues = { status: 'unexpected' } as Record; + console.log('Before validation:', parsedValues); + await assertion.fn(parsedValues); + console.log('After validation (defaults + catch applied):', parsedValues); + + const invalid = { status: 'ok', attempts: 10 }; + const failure = await assertion.fn(invalid); + console.log('Invalid parse result:', failure); +} + +void main(); diff --git a/src/examples/zod-signature-roundtrip.ts b/src/examples/zod-signature-roundtrip.ts new file mode 100644 index 00000000..786b5d4b --- /dev/null +++ b/src/examples/zod-signature-roundtrip.ts @@ -0,0 +1,38 @@ +import { z } from 'zod'; + +import { AxSignature } from '../ax/dsp/sig.js'; +import { getZodMetadata } from '../ax/zod/metadata.js'; + +const ticketSchema = z.object({ + id: z.string().uuid(), + severity: z.enum(['low', 'medium', 'high']).catch('low'), + summary: z.string().min(10), + tags: z.array(z.string()).default([]), +}); + +const ticketSignature = AxSignature.fromZod(ticketSchema, { + mode: 'safeParse', +}); + +console.log('Ticket outputs:', ticketSignature.getOutputFields()); + +const ticketMetadata = getZodMetadata(ticketSignature); +if (ticketMetadata) { + console.log('Ticket conversion issues:', ticketMetadata.issues); + console.log( + 'Round-trip preserves schema reference:', + ticketSignature.toZod() === ticketSchema + ); +} + +const unionSchema = z.object({ + payload: z.union([z.string(), z.number()]), +}); + +const unionSignature = AxSignature.fromZod(unionSchema); +console.log('Union outputs:', unionSignature.getOutputFields()); + +const unionMetadata = getZodMetadata(unionSignature); +if (unionMetadata) { + console.log('Union conversion issues:', unionMetadata.issues); +} From 8e87e2a4ebb8fc109ceaec251f35e709c124cf21 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Sat, 18 Oct 2025 19:55:16 +0700 Subject: [PATCH 2/3] feat(zod): enable schemas in ax entrypoints --- src/ax/dsp/generate.test.ts | 43 ++++++++++++++++++++++++++++++++++++ src/ax/dsp/generate.ts | 8 +++++-- src/ax/dsp/sig.test.ts | 22 +++++++++++++++++++ src/ax/dsp/template.ts | 44 ++++++++++++++++++++++++++++++------- src/ax/dsp/types.ts | 3 +++ src/ax/zod/convert.ts | 9 +++++++- 6 files changed, 118 insertions(+), 11 deletions(-) diff --git a/src/ax/dsp/generate.test.ts b/src/ax/dsp/generate.test.ts index 130ae1a2..f0864472 100644 --- a/src/ax/dsp/generate.test.ts +++ b/src/ax/dsp/generate.test.ts @@ -10,6 +10,8 @@ import type { AxChatResponse } from '../ai/types.js'; import { AxStopFunctionCallException } from './functions.js'; import { AxGen, type AxGenerateError } from './generate.js'; import { AxSignature } from './sig.js'; +import { getZodMetadata } from '../zod/metadata.js'; +import { ax } from './template.js'; import type { AxProgramForwardOptions } from './types.js'; function createStreamingResponse( @@ -732,6 +734,47 @@ describe('AxGen Message Validation', () => { }); describe('AxGen with Zod signatures', () => { + it('constructs directly from a Zod schema with custom options', () => { + const schema = z.object({ + title: z.string().min(3), + tags: z.array(z.string()).default([]), + }); + + const gen = new AxGen(schema, { + zod: { + assertionLevel: 'none', + mode: 'parse', + }, + }); + + const signature = gen.getSignature(); + const metadata = getZodMetadata(signature); + expect(metadata?.options.mode).toBe('parse'); + expect(metadata?.options.assertionLevel).toBe('none'); + const asserts = (gen as any).asserts as unknown[]; + expect(asserts?.length ?? 0).toBe(0); + }); + + it('creates a generator via ax() helper when given a Zod schema', () => { + const schema = z.object({ + topic: z.string(), + summary: z.string().catch('summary unavailable'), + }); + + const gen = ax(schema, { + zod: { + assertionLevel: 'final', + mode: 'safeParse', + }, + }); + + expect(gen).toBeInstanceOf(AxGen); + const asserts = (gen as any).asserts as unknown[]; + expect(asserts.length).toBeGreaterThan(0); + const metadata = getZodMetadata(gen.getSignature()); + expect(metadata?.schema).toBe(schema); + }); + it('applies final Zod assertions and default values', async () => { const schema = z.object({ name: z.string().min(1), diff --git a/src/ax/dsp/generate.ts b/src/ax/dsp/generate.ts index 13e18765..f6f72797 100644 --- a/src/ax/dsp/generate.ts +++ b/src/ax/dsp/generate.ts @@ -1,5 +1,7 @@ // ReadableStream is available globally in modern browsers and Node.js 16+ +import type { ZodTypeAny } from 'zod'; + import { type Context, context, @@ -141,13 +143,15 @@ export class AxGen constructor( signature: - | NonNullable[0]> - | AxSignature, + | ConstructorParameters[0] + | AxSignature + | ZodTypeAny, options?: Readonly> ) { super(signature, { description: options?.description, traceLabel: options?.traceLabel, + zod: options?.zod, }); this.options = options; diff --git a/src/ax/dsp/sig.test.ts b/src/ax/dsp/sig.test.ts index 0deb8408..f7d15f8b 100644 --- a/src/ax/dsp/sig.test.ts +++ b/src/ax/dsp/sig.test.ts @@ -965,6 +965,28 @@ describe('Zod integration', () => { ); }); + it('preserves array element types when inner schemas are optional', () => { + const schema = z.object({ + tags: z.array(z.string().optional()), + }); + + const signature = AxSignature.fromZod(schema); + const tagsField = signature + .getOutputFields() + .find((field) => field.name === 'tags'); + + expect(tagsField?.type?.name).toBe('string'); + expect(tagsField?.type?.isArray).toBe(true); + + const metadata = getZodMetadata(signature); + const hasArrayDowngrade = metadata?.issues.some( + (issue) => + issue.kind === 'unsupported' && + issue.path.toLowerCase().includes('tags') + ); + expect(hasArrayDowngrade).toBe(false); + }); + it('tracks downgrade telemetry for unions and maps them to json outputs', () => { const schema = z.object({ payload: z.union([z.string(), z.number()]), diff --git a/src/ax/dsp/template.ts b/src/ax/dsp/template.ts index 9bb350c9..cdbbe5fb 100644 --- a/src/ax/dsp/template.ts +++ b/src/ax/dsp/template.ts @@ -1,7 +1,10 @@ +import type { ZodTypeAny } from 'zod'; + import { AxGen } from './generate.js'; import { AxSignature } from './sig.js'; import type { ParseSignature } from './sigtypes.js'; import type { AxProgramForwardOptions } from './types.js'; +import { isZodSchema } from '../zod/util.js'; // Function for string-based type-safe signature creation export function s( @@ -10,7 +13,7 @@ export function s( return AxSignature.create(signature); } -// Function for type-safe generator creation - supports both strings and AxSignature objects +// Function for type-safe generator creation - supports strings, signatures, and Zod schemas export function ax< const T extends string, ThoughtKey extends string = 'thought', @@ -43,18 +46,37 @@ export function ax< : { [P in ThoughtKey]?: string }) >; export function ax< - T extends string | AxSignature, + TSchema extends ZodTypeAny, + ThoughtKey extends string = 'thought', +>( + signature: TSchema, + options?: Readonly< + AxProgramForwardOptions & { thoughtFieldName?: ThoughtKey } + > +): AxGen< + Record, + Record & + (string extends ThoughtKey + ? { thought?: string } + : { [P in ThoughtKey]?: string }) +>; +export function ax< + T extends string | AxSignature | ZodTypeAny, ThoughtKey extends string = 'thought', TInput extends Record = T extends string ? ParseSignature['inputs'] : T extends AxSignature ? I - : never, + : T extends ZodTypeAny + ? Record + : never, TOutput extends Record = T extends string ? ParseSignature['outputs'] : T extends AxSignature ? O - : never, + : T extends ZodTypeAny + ? Record + : never, >( signature: T, options?: Readonly< @@ -67,10 +89,16 @@ export function ax< ? { thought?: string } : { [P in ThoughtKey]?: string }) > { - const typedSignature = - typeof signature === 'string' - ? AxSignature.create(signature) - : (signature as AxSignature); + const typedSignature = (() => { + if (typeof signature === 'string') { + return AxSignature.create(signature); + } + if (isZodSchema(signature)) { + return signature; + } + return signature as AxSignature; + })(); + return new AxGen< TInput, TOutput & diff --git a/src/ax/dsp/types.ts b/src/ax/dsp/types.ts index 423f43eb..7fcb403d 100644 --- a/src/ax/dsp/types.ts +++ b/src/ax/dsp/types.ts @@ -125,6 +125,9 @@ export type AxProgramForwardOptions = AxAIServiceOptions & { // Field prefix is required for single output field programs strictMode?: boolean; + + // Schema conversions + zod?: AxZodSignatureOptions; }; export type AxAIServiceActionOptions< diff --git a/src/ax/zod/convert.ts b/src/ax/zod/convert.ts index c54c51b9..aaa0277f 100644 --- a/src/ax/zod/convert.ts +++ b/src/ax/zod/convert.ts @@ -188,8 +188,15 @@ const mapZodToAxFieldType = ( issues: AxZodConversionIssue[] ): Omit => { if (schema instanceof ZodArray) { + const elementPath = `${path}[]`; + const unwrapped = unwrapSchema(schema.element, elementPath); + issues.push(...unwrapped.notes); const innerIssues: AxZodConversionIssue[] = []; - const inner = mapZodToAxFieldType(schema.element, `${path}[]`, innerIssues); + const inner = mapZodToAxFieldType( + unwrapped.schema, + elementPath, + innerIssues + ); issues.push(...innerIssues); return { ...inner, From 663a0e92be239816a8a73f3a0aad749310a7a302 Mon Sep 17 00:00:00 2001 From: Tom X Nguyen Date: Sat, 18 Oct 2025 20:11:17 +0700 Subject: [PATCH 3/3] docs(zod): clean up docs --- docs/ZOD_INTEGRATION.md | 125 ++++++++++++++++++---------------------- 1 file changed, 56 insertions(+), 69 deletions(-) diff --git a/docs/ZOD_INTEGRATION.md b/docs/ZOD_INTEGRATION.md index 8b690b39..d439b52a 100644 --- a/docs/ZOD_INTEGRATION.md +++ b/docs/ZOD_INTEGRATION.md @@ -1,25 +1,28 @@ # Zod Integration Blueprint for Ax ## Overview + This document proposes a deep, foundational integration of [Zod](https://zod.dev) into Ax so that Zod schemas become first-class citizens throughout the signature, engine, and assertion layers. The design builds upon the groundwork in PR #388 and aims to unlock seamless schema reuse, rich runtime validation, and best-in-class developer ergonomics for AI workflow authoring. --- ## Architecture Proposal + The integration introduces new adapters and validation flows that allow Zod schemas to travel with Ax signatures from definition to runtime enforcement. ### Component Diagram (textual UML) + ``` -+-----------------+ +--------------------+ +---------------------+ ++-----------------+ +--------------------+ +---------------------+ | Zod Schema | --(register)--> | AxZodRegistry | --(build)-> | AxSignatureFactory | -| (z.object(...)) | | (schema cache) | | (existing) | -+-----------------+ +--------------------+ +---------------------+ +| (z.object(...)) | | (schema cache) | | (existing) | ++-----------------+ +--------------------+ +---------------------+ | | | | v v - | +----------------+ +--------------------+ + | +----------------+ +---------------------+ | | AxSignature |<--+ | AxGen / AxFlow | | | (with zodMeta) | | | (DSPy-style engine) | - | +----------------+ | +--------------------+ + | +----------------+ | +---------------------+ | | | | | v | v | +----------------+ | +----------------------+ @@ -28,10 +31,10 @@ The integration introduces new adapters and validation flows that allow Zod sche | +----------------+ +----------------------+ | | | | v v - +-----------------------------+ +--------------+ +-------------------------+ - | Runtime Output | --LLM--> | Streaming Validator | - | (JSON/prose) | | (per-field / final) | - +--------------+-+ +-------------------------+ + | +----------------+ +-------------------------+ + +------------------------------+ | Runtime Output | ---LLM--> | Streaming Validator | + | (JSON/prose) | | (per-field / final) | + +----------------+ +-------------------------+ | v +-------------------+ @@ -40,31 +43,39 @@ The integration introduces new adapters and validation flows that allow Zod sche | - errors | | - telemetry | +-------------------+ + ``` ### Key Modules & APIs + - `AxSignature.fromZod(schema, options?: AxZodSignatureOptions): AxSignature` + - `options.strict`: throw on any downgrade or unsupported feature. - `options.streaming`: enable emission of field-level validators for streaming output enforcement. - `options.mode`: `'parse' | 'safeParse' | 'coerce'` to align with Zod parsing semantics. - `options.assertionLevel`: `'none' | 'final' | 'streaming' | 'both'` to control auto-assertion wiring. - `AxSignature.toZod(signature, options?: AxToZodOptions): ZodSchema` + - Round-trip support with metadata preservation. - `AxZodRegistry` + - Internal cache keyed by signature ID to store original Zod schema, downgrade notes, and validation options. - Provides `get(schemaId)` for runtime validation modules. - `ZodAssertionAdapter` + - Translates Zod `.parse`, `.safeParse`, `.min`, `.max`, `.default`, `.catch`, `.transform`, and `.refine` hooks into Ax assertions. - Emits `AxAssertion` objects with `severity`, `recovery` (fallback result), and `telemetry` payloads. - `StreamingZodValidator` + - Wraps Zod schema introspection to produce field-level validators suitable for Ax’s streaming extraction pipeline. - Supports chunk-level validation and progressive parse with buffering. - `AxValidationTelemetry` + - Unified event schema capturing downgrade issues, parse failures, defaults applied, and user-facing remediation tips. - `AxZodCLI` @@ -72,103 +83,85 @@ The integration introduces new adapters and validation flows that allow Zod sche --- -## Implementation Roadmap - -### Phase 1 – Schema Fidelity & Runtime Parse (Foundational) -1. **Metadata Extensions**: Augment `AxSignature` to carry `zodMeta` (original schema, version, options) via weak references to avoid bundling heavy schema graphs when unused. -2. **Round-trip APIs**: Harden `AxSignature.fromZod` and add `AxSignature.toZod`, targeting >95% fidelity for primitives, objects, literals, enums, arrays, unions, records, defaults, transforms, `ZodEffects`, and branded schemas. -3. **Assertion Wiring**: Auto-register final-result assertions by default using `.safeParse`. Surface `.error.flatten()` payloads in telemetry. -4. **Downgrade/Validation Telemetry**: Extend existing conversion diagnostics with runtime parse failure logging, `strict` mode enforcement, and actionable guidance. -5. **Deterministic Tests**: Add comprehensive unit tests that cover round-tripping, default application, safe parsing, and error telemetry without external API calls. -6. **Docs & Examples**: Update README, SIGNATURES guide, and create a dedicated `ZOD_INTEGRATION` doc (this file) plus deterministic examples (e.g., using mocked LLM outputs). - -### Phase 2 – Streaming Validation & Advanced Semantics -1. **Streaming Validator**: Design `StreamingZodValidator` to map object fields to incremental validators. Support buffering for arrays/objects and fail-fast on impossible states. -2. **Per-field Assertions**: Augment AxGen to attach streaming assertions that fire as tokens arrive, bridging Ax’s internal validator with Zod hints (min/max length, regex, `.nonempty`, etc.). -3. **Fallback Strategies**: Integrate jsonrepair/normalization hooks before final parse; allow `.catch` and `.default` to recover gracefully. -4. **Transforms & Effects**: Provide opt-in handling for `ZodEffects` and `transform` pipelines by running transforms post-parse while tracking original raw output for telemetry. -5. **CLI Enhancements**: Extend `ax zod audit` to simulate streaming validations and suggest signature adjustments. - -### Phase 3 – Ecosystem & Extensibility -1. **Factory & Recipes**: Provide helper factories (`createZodSignature(schema, options)`) and recipe docs for migrating from Mastra/Superstruct flows. -2. **Extensibility Hooks**: Formalize plugin interface so alternative schema libraries (Superstruct, Valibot) can plug into the same assertion pipeline. -3. **Performance Tuning**: Benchmark large schemas (>30 fields) under streaming and batch parse; expose profiling telemetry. -4. **DX Enhancements**: VS Code snippets, typed helper utilities, and cross-linked guides (e.g., flow recipes, optimization docs). -5. **Community Feedback Loop**: Add telemetry counters (opt-in) for parse failures and defaults applied; document contribution guidelines for new schema adapters. - ---- - ## Code Snippets ### 1. Converting Zod Schema to Ax Signature with Strict Validation + ```ts -import { z } from "zod"; -import { AxSignature } from "ax/signature"; +import { z } from 'zod' +import { AxSignature } from 'ax/signature' const invoiceSchema = z.object({ invoiceId: z.string().uuid(), totalCents: z.number().int().min(0), - issuedAt: z.string().datetime().default(() => new Date().toISOString()), -}); + issuedAt: z + .string() + .datetime() + .default(() => new Date().toISOString()), +}) const signature = AxSignature.fromZod(invoiceSchema, { strict: true, - assertionLevel: "final", - mode: "safeParse", -}); + assertionLevel: 'final', + mode: 'safeParse', +}) ``` ### 2. Using Auto-Applied Zod Assertions in AxGen + ```ts -import { axgen } from "ax"; -import { z } from "zod"; +import { axgen } from 'ax' +import { z } from 'zod' const schema = z.object({ customerName: z.string().min(1), - preferredContact: z.enum(["email", "phone"]).default("email"), -}); + preferredContact: z.enum(['email', 'phone']).default('email'), +}) const gen = axgen() .signature(schema) // equivalent to AxSignature.fromZod(schema) - .prompt("Collect the customer preferences from the conversation.") + .prompt('Collect the customer preferences from the conversation.') .onFailure((ctx, error) => { - ctx.logger.warn("Zod validation failed", { error }); - return ctx.retry(); - }); + ctx.logger.warn('Zod validation failed', { error }) + return ctx.retry() + }) -const result = await gen.run({ transcript }); +const result = await gen.run({ transcript }) // result.payload is guaranteed to satisfy schema.parse(...) semantics ``` ### 3. Streaming Validation Hook + ```ts -import { createStreamingValidator } from "ax/zod/stream"; +import { createStreamingValidator } from 'ax/zod/stream' -const validator = createStreamingValidator(schema, { chunkSize: 128 }); +const validator = createStreamingValidator(schema, { chunkSize: 128 }) for await (const token of llmStream) { - const status = validator.ingest(token); - if (status.type === "error") { + const status = validator.ingest(token) + if (status.type === 'error') { // Apply jsonrepair, request model correction, or abort - await controller.requestFix(status.issues); + await controller.requestFix(status.issues) } } -const finalValue = validator.finalize(); +const finalValue = validator.finalize() ``` ### 4. Round-trip Conversion (AxSignature → Zod) + ```ts -const existingSignature = AxSignature.load("customer.profile"); -const zodSchema = existingSignature.toZod(); +const existingSignature = AxSignature.load('customer.profile') +const zodSchema = existingSignature.toZod() // Apply additional refinements with Zod API const stricterSchema = zodSchema.extend({ - loyaltyTier: z.enum(["bronze", "silver", "gold"]).catch("bronze"), -}); + loyaltyTier: z.enum(['bronze', 'silver', 'gold']).catch('bronze'), +}) ``` ### 5. CLI Audit Example + ```bash npx ax zod audit ./schemas/customer.ts --strict --report json ``` @@ -176,6 +169,7 @@ npx ax zod audit ./schemas/customer.ts --strict --report json --- ## Trade-offs & Risks + - **Performance Overhead**: Zod parsing incurs runtime cost, especially for large schemas. Mitigation: allow opt-in streaming validation to catch errors early, cache compiled schemas, and benchmark to ensure <10% throughput regression for typical flows. - **Streaming Complexity**: Zod is not inherently streaming-aware. The proposed streaming validator will require custom buffering logic; certain constructs (e.g., regex on entire strings, cross-field refinement) may only be enforceable at finalize-time. - **Schema Fidelity Gaps**: Features like `ZodFunction`, complex intersections, and advanced `ZodEffects` may still need JSON fallbacks. Document these limitations and provide telemetry to guide users. @@ -185,16 +179,9 @@ npx ax zod audit ./schemas/customer.ts --strict --report json --- ## Success Metrics + - **Schema Fidelity**: ≥95% of schemas in the test corpus convert round-trip without downgrade warnings. - **Runtime Reliability**: ≥90% reduction in manual signature authoring among beta users migrating from Zod-heavy stacks (survey-based). - **DX Satisfaction**: Positive qualitative feedback (≥4/5) from developer surveys on the new CLI, docs, and auto-assertions. - **Telemetry Signals**: Monitoring shows decreasing parse failure rates over time due to better defaults and `.catch` handling; streaming validator issues are actionable. - **Adoption**: At least two reference integrations (e.g., Mastra migration, OpenAI structured extraction recipe) published using the new APIs. - ---- - -## Next Steps -- Circulate this blueprint with maintainers (@dosco, @monotykamary) for feedback. -- Align roadmap milestones with upcoming releases. -- Kick off Phase 1 implementation with focus on metadata extensions and runtime assertions. -