Skip to content

Commit 6b8cc14

Browse files
lgrammellgrammel
authored
feat (ai/core): support recursive zod schemas (#4702)
1 parent 9c605ca commit 6b8cc14

File tree

8 files changed

+236
-7
lines changed

8 files changed

+236
-7
lines changed

‎.changeset/rude-flies-own.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@ai-sdk/ui-utils': patch
3+
---
4+
5+
feat (ai/core): support recursive zod schemas

‎content/docs/02-foundations/04-tools.mdx

+5-3
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,11 @@ using [multi-step calls](/docs/ai-sdk-core/tools-and-tool-calling#multi-step-cal
4343

4444
Schemas are used to define the parameters for tools and to validate the [tool calls](/docs/ai-sdk-core/tools-and-tool-calling).
4545

46-
The AI SDK supports both raw JSON schemas (using the `jsonSchema` function) and [Zod](https://zod.dev/) schemas.
47-
[Zod](https://zod.dev/) is the most popular JavaScript schema validation library.
48-
You can install Zod with:
46+
The AI SDK supports both raw JSON schemas (using the [`jsonSchema` function](/docs/reference/ai-sdk-core/json-schema))
47+
and [Zod](https://zod.dev/) schemas (either directly or using the [`zodSchema` function](/docs/reference/ai-sdk-core/zod-schema)).
48+
49+
[Zod](https://zod.dev/) is a popular TypeScript schema validation library.
50+
You can install it with:
4951

5052
<Tabs items={['pnpm', 'npm', 'yarn']}>
5153
<Tab>

‎content/docs/03-ai-sdk-core/10-generating-structured-data.mdx

+6-1
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,14 @@ with the [`generateObject`](/docs/reference/ai-sdk-core/generate-object)
1616
and [`streamObject`](/docs/reference/ai-sdk-core/stream-object) functions.
1717
You can use both functions with different output strategies, e.g. `array`, `object`, or `no-schema`,
1818
and with different generation modes, e.g. `auto`, `tool`, or `json`.
19-
You can use [Zod schemas](./schemas-and-zod) or [JSON schemas](/docs/reference/ai-sdk-core/json-schema) to specify the shape of the data that you want,
19+
You can use [Zod schemas](/docs/reference/ai-sdk-core/zod-schema) or [JSON schemas](/docs/reference/ai-sdk-core/json-schema) to specify the shape of the data that you want,
2020
and the AI model will generate data that conforms to that structure.
2121

22+
<Note>
23+
You can pass Zod objects directly to the AI SDK functions or use the
24+
`zodSchema` helper function.
25+
</Note>
26+
2227
## Generate Object
2328

2429
The `generateObject` generates structured data from a prompt.

‎content/docs/07-reference/01-ai-sdk-core/26-zod-schema.mdx

+89
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
---
2+
title: zodSchema
3+
description: Helper function for creating Zod schemas
4+
---
5+
6+
# `zodSchema()`
7+
8+
`zodSchema` is a helper function that converts a Zod schema into a JSON schema object that is compatible with the AI SDK.
9+
It takes a Zod schema and optional configuration as inputs, and returns a typed schema.
10+
11+
You can use it to [generate structured data](/docs/ai-sdk-core/generating-structured-data) and in [tools](/docs/ai-sdk-core/tools-and-tool-calling).
12+
13+
<Note>
14+
You can also pass Zod objects directly to the AI SDK functions. Internally,
15+
the AI SDK will convert the Zod schema to a JSON schema using `zodSchema()`.
16+
However, if you want to specify options such as `useReferences`, you can pass
17+
the `zodSchema()` helper function instead.
18+
</Note>
19+
20+
## Example with recursive schemas
21+
22+
```ts
23+
import { zodSchema } from 'ai';
24+
import { z } from 'zod';
25+
26+
// Define a base category schema
27+
const baseCategorySchema = z.object({
28+
name: z.string(),
29+
});
30+
31+
// Define the recursive Category type
32+
type Category = z.infer<typeof baseCategorySchema> & {
33+
subcategories: Category[];
34+
};
35+
36+
// Create the recursive schema using z.lazy
37+
const categorySchema: z.ZodType<Category> = baseCategorySchema.extend({
38+
subcategories: z.lazy(() => categorySchema.array()),
39+
});
40+
41+
// Create the final schema with useReferences enabled for recursive support
42+
const mySchema = zodSchema(
43+
z.object({
44+
category: categorySchema,
45+
}),
46+
{ useReferences: true },
47+
);
48+
```
49+
50+
## Import
51+
52+
<Snippet text={`import { zodSchema } from "ai"`} prompt={false} />
53+
54+
## API Signature
55+
56+
### Parameters
57+
58+
<PropertiesTable
59+
content={[
60+
{
61+
name: 'zodSchema',
62+
type: 'z.Schema',
63+
description: 'The Zod schema definition.',
64+
},
65+
{
66+
name: 'options',
67+
type: 'object',
68+
description: 'Additional options for the schema conversion.',
69+
properties: [
70+
{
71+
type: 'object',
72+
parameters: [
73+
{
74+
name: 'useReferences',
75+
isOptional: true,
76+
type: 'boolean',
77+
description:
78+
'Enables support for references in the schema. This is required for recursive schemas, e.g. with `z.lazy`. However, not all language models and providers support such references. Defaults to `false`.',
79+
},
80+
],
81+
},
82+
],
83+
},
84+
]}
85+
/>
86+
87+
### Returns
88+
89+
A Schema object that is compatible with the AI SDK, containing both the JSON schema representation and validation functionality.

‎content/docs/07-reference/01-ai-sdk-core/index.mdx

+5
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,11 @@ It also contains the following helper functions:
6969
description: 'Creates AI SDK compatible JSON schema objects.',
7070
href: '/docs/reference/ai-sdk-core/json-schema',
7171
},
72+
{
73+
title: 'zodSchema()',
74+
description: 'Creates AI SDK compatible Zod schema objects.',
75+
href: '/docs/reference/ai-sdk-core/zod-schema',
76+
},
7277
{
7378
title: 'createProviderRegistry()',
7479
description:

‎packages/ui-utils/src/__snapshots__/zod-schema.test.ts.snap

+72-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ exports[`zodSchema > json schema conversion > should create a schema with simple
2020
}
2121
`;
2222

23-
exports[`zodSchema > json schema conversion > should duplicate referenced schemas (and not use references) 1`] = `
23+
exports[`zodSchema > json schema conversion > should duplicate referenced schemas (and not use references) by default 1`] = `
2424
{
2525
"$schema": "http://json-schema.org/draft-07/schema#",
2626
"additionalProperties": false,
@@ -185,3 +185,74 @@ exports[`zodSchema > json schema conversion > should support required enums 1`]
185185
"type": "object",
186186
}
187187
`;
188+
189+
exports[`zodSchema > json schema conversion > should use recursive references with z.lazy when useReferences is true 1`] = `
190+
{
191+
"$schema": "http://json-schema.org/draft-07/schema#",
192+
"additionalProperties": false,
193+
"properties": {
194+
"category": {
195+
"additionalProperties": false,
196+
"properties": {
197+
"name": {
198+
"type": "string",
199+
},
200+
"subcategories": {
201+
"items": {
202+
"$ref": "#/properties/category",
203+
},
204+
"type": "array",
205+
},
206+
},
207+
"required": [
208+
"name",
209+
"subcategories",
210+
],
211+
"type": "object",
212+
},
213+
},
214+
"required": [
215+
"category",
216+
],
217+
"type": "object",
218+
}
219+
`;
220+
221+
exports[`zodSchema > json schema conversion > should use references when useReferences is true 1`] = `
222+
{
223+
"$schema": "http://json-schema.org/draft-07/schema#",
224+
"additionalProperties": false,
225+
"properties": {
226+
"group1": {
227+
"items": {
228+
"additionalProperties": false,
229+
"properties": {
230+
"number": {
231+
"type": "number",
232+
},
233+
"text": {
234+
"type": "string",
235+
},
236+
},
237+
"required": [
238+
"text",
239+
"number",
240+
],
241+
"type": "object",
242+
},
243+
"type": "array",
244+
},
245+
"group2": {
246+
"items": {
247+
"$ref": "#/properties/group1/items",
248+
},
249+
"type": "array",
250+
},
251+
},
252+
"required": [
253+
"group1",
254+
"group2",
255+
],
256+
"type": "object",
257+
}
258+
`;

‎packages/ui-utils/src/zod-schema.test.ts

+41-1
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ describe('zodSchema', () => {
7777
expect(schema.jsonSchema).toMatchSnapshot();
7878
});
7979

80-
it('should duplicate referenced schemas (and not use references)', () => {
80+
it('should duplicate referenced schemas (and not use references) by default', () => {
8181
const Inner = z.object({
8282
text: z.string(),
8383
number: z.number(),
@@ -92,6 +92,46 @@ describe('zodSchema', () => {
9292

9393
expect(schema.jsonSchema).toMatchSnapshot();
9494
});
95+
96+
it('should use references when useReferences is true', () => {
97+
const Inner = z.object({
98+
text: z.string(),
99+
number: z.number(),
100+
});
101+
102+
const schema = zodSchema(
103+
z.object({
104+
group1: z.array(Inner),
105+
group2: z.array(Inner),
106+
}),
107+
{ useReferences: true },
108+
);
109+
110+
expect(schema.jsonSchema).toMatchSnapshot();
111+
});
112+
113+
it('should use recursive references with z.lazy when useReferences is true', () => {
114+
const baseCategorySchema = z.object({
115+
name: z.string(),
116+
});
117+
118+
type Category = z.infer<typeof baseCategorySchema> & {
119+
subcategories: Category[];
120+
};
121+
122+
const categorySchema: z.ZodType<Category> = baseCategorySchema.extend({
123+
subcategories: z.lazy(() => categorySchema.array()),
124+
});
125+
126+
const schema = zodSchema(
127+
z.object({
128+
category: categorySchema,
129+
}),
130+
{ useReferences: true },
131+
);
132+
133+
expect(schema.jsonSchema).toMatchSnapshot();
134+
});
95135
});
96136

97137
describe('output validation', () => {

‎packages/ui-utils/src/zod-schema.ts

+13-1
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,22 @@ import { jsonSchema, Schema } from './schema';
55

66
export function zodSchema<OBJECT>(
77
zodSchema: z.Schema<OBJECT, z.ZodTypeDef, any>,
8+
options?: {
9+
/**
10+
* Enables support for references in the schema.
11+
* This is required for recursive schemas, e.g. with `z.lazy`.
12+
* However, not all language models and providers support such references.
13+
* Defaults to `false`.
14+
*/
15+
useReferences?: boolean;
16+
},
817
): Schema<OBJECT> {
18+
// default to no references (to support openapi conversion for google)
19+
const useReferences = options?.useReferences ?? false;
20+
921
return jsonSchema(
1022
zodToJsonSchema(zodSchema, {
11-
$refStrategy: 'none', // no references (to support openapi conversion for google)
23+
$refStrategy: useReferences ? 'root' : 'none',
1224
target: 'jsonSchema7', // note: openai mode breaks various gemini conversions
1325
}) as JSONSchema7,
1426
{

0 commit comments

Comments
 (0)