Skip to content

Conversation

monotykamary
Copy link
Contributor

@monotykamary monotykamary commented Oct 15, 2025

  • What kind of change does this PR introduce? (Bug fix, feature, docs update, ...)
    Feature + tooling polish: end-to-end Zod schema support, downgrade diagnostics, docs/examples, and several conversion bug fixes.

  • What is the current behavior? (You can also link to an open issue here)
    Ax signatures must be authored manually and existing Zod schemas can’t be reused. Any fallbacks (records/maps/unions/transforms) happen silently, array optionality leaks from element wrappers, and there’s no discoverable guidance on Zod → Ax mapping. Issue Standard Schema for Signature Input/Output? #348 captures the gap.

  • What is the new behavior (if this is a feature change)?
    AxSignature.fromZod converts Zod input/output objects (v3/v4) into signatures with full type inference and downgrade telemetry.
    AxSignature.debugZodConversion, getZodConversionIssues, and reportZodConversionIssues expose conversion metadata; strict mode now throws with actionable suggestions.
    • Conversion logic handles primitives, literals/enums/native enums, branded schemas, nested arrays, discriminated unions, records/maps, and now explicitly flags effects/pipe transformations so runtime-altering schemas don’t masquerade as primitives.
    • Fixed optional-element arrays leaking optionality, stabilized literal/class downgrade handling, and re-exported the new AxSignatureFromZodOptions.
    • Docs (README.md, docs/SIGNATURES.md) outline the mapping matrix, downgrade semantics, and auditing workflow.
    • Added src/examples/zod-signature-example.ts (console demo) and src/examples/zod-live-example.ts (OpenAI-backed flow + AxGen using AxSignature.fromZod).
    • Expanded src/ax/dsp/zodToSignature.test.ts with comprehensive coverage (supported scenarios, downgrade cases, branded schemas, nested structures, transforms).
    • Declared a zod peer dependency so consumers get a compatible version.

  • Other information:
    • Effects/pipes are conservatively downgraded to json; if future work unlocks richer Ax types for transformed schemas we can adjust.
    • Zod constructs that Ax can’t represent yet (e.g., ZodFunction, complex intersections) intentionally fall back to json; conversion telemetry makes these visible but they remain unsupported.
    • Example coverage relies on an OPENAI_APIKEY; adding a deterministic fixture-driven smoke test in a follow-up issue would harden CI coverage.
    • Docs now surface the feature, but discoverability across other guides (e.g., flow recipes) might benefit from cross-links in a future polish pass.


npm run tsx -- src/examples/zod-live-example.ts
$ npm run tsx -- src/examples/zod-live-example.ts

> @ax-llm/ax-monorepo@14.0.31 tsx
> node --env-file=.env --import=tsx src/examples/zod-live-example.ts

[zod-live-example] ticket signature:
"Classify incoming incidents and propose a triage plan" summary:string "Short description of the incident being triaged", details:string "Full report including any reproduction steps", severity:string, tags?:string[] "Optional labels shared with observability dashboards" -> recommendedAction:class "escalate | monitor | close", priority:class "P0 | P1 | P2 | P3", suggestedOwners?:string[]

[zod-live-example] running AxGen.fromZod direct example...
[zod-live-example] direct result:
{
  recommendedAction: 'escalate',
  priority: 'P0',
  suggestedOwners: [ 'checkout-team', 'payments-team', 'oncall-sre' ]
}
[AxFlow] new AxFlow() is deprecated. Use flow() factory instead.

[zod-live-example] running AxFlow + fromZod example...
[zod-live-example] flow result:
{
  recommendedAction: 'escalate',
  priority: 'P0',
  suggestedOwners: [ 'checkout-service', 'payments-team', 'sre-oncall' ]

npm run tsx -- src/examples/zod-signature-example.ts
$ npm run tsx -- src/examples/zod-signature-example.ts

> @ax-llm/ax-monorepo@14.0.31 tsx
> node --env-file=.env --import=tsx src/examples/zod-signature-example.ts

Bug report signature inputs:
┌─────────┬──────────────┬──────────────────────────────────────┬──────────────────────────────────────────────────┬───────────────┬────────────┐
│ (index) │ name         │ type                                 │ description                                      │ title         │ isOptional │
├─────────┼──────────────┼──────────────────────────────────────┼──────────────────────────────────────────────────┼───────────────┼────────────┤
│ 0       │ 'summary'    │ { name: 'string' }                   │ 'Short title for the issue''Summary'     │            │
│ 1       │ 'details'    │ { name: 'string' }                   │ 'Long form description supplied by the reporter''Details'     │            │
│ 2       │ 'severity'   │ { name: 'string', options: [Array] } │                                                  │ 'Severity'true       │
│ 3       │ 'labels'     │ [Object]                             │                                                  │ 'Labels'true       │
│ 4       │ 'reportedAt' │ { name: 'date' }                     │                                                  │ 'Reported At' │            │
└─────────┴──────────────┴──────────────────────────────────────┴──────────────────────────────────────────────────┴───────────────┴────────────┘

Bug report signature outputs:
┌─────────┬─────────────────────┬─────────────────────────────────────┬──────────────────────┐
│ (index) │ name                │ type                                │ title                │
├─────────┼─────────────────────┼─────────────────────────────────────┼──────────────────────┤
│ 0       │ 'triageSummary'     │ { name: 'string' }                  │ 'Triage Summary'     │
│ 1       │ 'suggestedPriority' │ { name: 'class', options: [Array] } │ 'Suggested Priority' │
│ 2       │ 'requiresHotfix'    │ { name: 'boolean' }                 │ 'Requires Hotfix'    │
└─────────┴─────────────────────┴─────────────────────────────────────┴──────────────────────┘

Round-tripped Zod schemas: {
  inputShape: [ 'summary', 'details', 'severity', 'labels', 'reportedAt' ],
  outputShape: [ 'triageSummary', 'suggestedPriority', 'requiresHotfix' ]
}

Sample request payload:
{
  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',
  labels: [ 'ui', 'api' ],
  reportedAt: 2025-10-16T03:40:22.481Z
}

Call `forward` with an AxAI instance to run the classification.

@monotykamary
Copy link
Contributor Author

err there might be a handful of cases that I might be overlooking, but hopefully it should be a good scaffold for supporting Zod.

@dosco
Copy link
Collaborator

dosco commented Oct 15, 2025

I've gotten a lot of feedback to support zod and structured extraction like enforcing schema's with json objects etc. While I'm also a user of zod people have pointed me to other libs like https://github.com/ianstormtaylor/superstruct and said we should support. I'm ok with being opinionated and just pick zod since it's the most popular.

@dosco
Copy link
Collaborator

dosco commented Oct 15, 2025

I guess the big question I have is whats the point of generating signatures from zod if zod validation (eg .min(10)) is not being used? We already have two apis for generating signatures the f() fluent api and the direct string line input -> output and both auto generate types.

Not against this just trying to understand the usecase better cause I know people will just say don't convert to signature (no downgrading) and just use zod validation for input's and outputs.

@monotykamary
Copy link
Contributor Author

Likewise, I agree. There are things like .catch/.default, min/max, and other rich validation that would require tackling a lot more on the sig engine side of Ax to make Zod a first-class citizen.

This PR would likely only cover stories such as if they are migrating from something like Mastra workflows (that is Zod heavy) and they wish to convert it to signatures rather than duplicate schemas. Although, I don't think I made a public round-trip API to convert signatures -> Zod, like DSPy signatures -> Pydantic yet. Let me add those in a bit just for reference.

If the use-case is more first-class and foundational support for Zod, that's probably where we begin to diverge from the simple DSPy and begin to piggyback some of the more niceties of Zod. .catch/.default, safeParse, and maybe helper libraries like in traversable/zod would help build the forgiving parts of the system, while giving us much more failsafes against malformed DSPy-like prompt outputs that happens more often around 30+ fields or if the description prompt tries to force a format.

I do have a sanitize -> jsonrepair, normalize -> try schema -> fallback -> final schema pipeline for Mastra workflows in another project, but that's really to get around models that can't do structured JSON well. The DSPy yaml-like outputs are much more forgiving and prose-friendly (which is why I'm here lol); adding foundational Zod support would be icing on top of this cake.

@dosco
Copy link
Collaborator

dosco commented Oct 16, 2025

we do have dspy assertions maybe the zod config should be stored in the signature and then zod.parse could be automatically added to the assertion list when the signature is used with an axgen see .addAssert(). and the parsed value returned.

the internal signature validator can still run for the streaming validation and the zod one on the final result. this would be even better if we can call the validator of each field in the zod object as the field become available. that way it can work with streaming. (i don't know if we can do that with zod)

@dosco
Copy link
Collaborator

dosco commented Oct 16, 2025

i'm open to deep integration with zod.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants