Skip to content

Commit 3c1224d

Browse files
RobertCraigieZijiaZhang
authored and
stainless-app[bot]
committed
fix(json-schema): correct handling of nested recursive schemas (#992)
* Fix zod to json schema with nested and recursive objects * minor style updates * add an iteration limit --------- Co-authored-by: Zijia Zhang <zijia.zhang99@gmail.com>
1 parent d486d27 commit 3c1224d

File tree

4 files changed

+242
-27
lines changed

4 files changed

+242
-27
lines changed

‎src/_vendor/zod-to-json-schema/Options.ts

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,9 @@ export type Options<Target extends Targets = 'jsonSchema7'> = {
3838
openaiStrictMode?: boolean;
3939
};
4040

41-
export const defaultOptions: Options = {
41+
const defaultOptions: Omit<Options, 'definitions' | 'basePath'> = {
4242
name: undefined,
4343
$refStrategy: 'root',
44-
basePath: ['#'],
4544
effectStrategy: 'input',
4645
pipeStrategy: 'all',
4746
dateStrategy: 'format:date-time',
@@ -51,7 +50,6 @@ export const defaultOptions: Options = {
5150
definitionPath: 'definitions',
5251
target: 'jsonSchema7',
5352
strictUnions: false,
54-
definitions: {},
5553
errorMessages: false,
5654
markdownDescription: false,
5755
patternStrategy: 'escape',
@@ -63,13 +61,20 @@ export const defaultOptions: Options = {
6361

6462
export const getDefaultOptions = <Target extends Targets>(
6563
options: Partial<Options<Target>> | string | undefined,
66-
) =>
67-
(typeof options === 'string' ?
68-
{
69-
...defaultOptions,
70-
name: options,
71-
}
72-
: {
73-
...defaultOptions,
74-
...options,
75-
}) as Options<Target>;
64+
) => {
65+
// We need to add `definitions` here as we may mutate it
66+
return (
67+
typeof options === 'string' ?
68+
{
69+
...defaultOptions,
70+
basePath: ['#'],
71+
definitions: {},
72+
name: options,
73+
}
74+
: {
75+
...defaultOptions,
76+
basePath: ['#'],
77+
definitions: {},
78+
...options,
79+
}) as Options<Target>;
80+
};

‎src/_vendor/zod-to-json-schema/zodToJsonSchema.ts

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -49,14 +49,28 @@ const zodToJsonSchema = <Target extends Targets = 'jsonSchema7'>(
4949
}
5050

5151
const definitions: Record<string, any> = {};
52+
const processedDefinitions = new Set();
5253

53-
for (const [name, zodSchema] of Object.entries(refs.definitions)) {
54-
definitions[name] =
55-
parseDef(
56-
zodDef(zodSchema),
57-
{ ...refs, currentPath: [...refs.basePath, refs.definitionPath, name] },
58-
true,
59-
) ?? {};
54+
// the call to `parseDef()` here might itself add more entries to `.definitions`
55+
// so we need to continually evaluate definitions until we've resolved all of them
56+
//
57+
// we have a generous iteration limit here to avoid blowing up the stack if there
58+
// are any bugs that would otherwise result in us iterating indefinitely
59+
for (let i = 0; i < 500; i++) {
60+
const newDefinitions = Object.entries(refs.definitions).filter(
61+
([key]) => !processedDefinitions.has(key),
62+
);
63+
if (newDefinitions.length === 0) break;
64+
65+
for (const [key, schema] of newDefinitions) {
66+
definitions[key] =
67+
parseDef(
68+
zodDef(schema),
69+
{ ...refs, currentPath: [...refs.basePath, refs.definitionPath, key] },
70+
true,
71+
) ?? {};
72+
processedDefinitions.add(key);
73+
}
6074
}
6175

6276
return definitions;

‎tests/lib/__snapshots__/parser.test.ts.snap

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,34 @@ exports[`.parse() zod nested schema extraction 2`] = `
8484
"
8585
`;
8686

87+
exports[`.parse() zod recursive schema extraction 2`] = `
88+
"{
89+
"id": "chatcmpl-9vdbw9dekyUSEsSKVQDhTxA2RCxcK",
90+
"object": "chat.completion",
91+
"created": 1723523988,
92+
"model": "gpt-4o-2024-08-06",
93+
"choices": [
94+
{
95+
"index": 0,
96+
"message": {
97+
"role": "assistant",
98+
"content": "{\\"linked_list\\":{\\"value\\":1,\\"next\\":{\\"value\\":2,\\"next\\":{\\"value\\":3,\\"next\\":{\\"value\\":4,\\"next\\":{\\"value\\":5,\\"next\\":null}}}}}}",
99+
"refusal": null
100+
},
101+
"logprobs": null,
102+
"finish_reason": "stop"
103+
}
104+
],
105+
"usage": {
106+
"prompt_tokens": 40,
107+
"completion_tokens": 38,
108+
"total_tokens": 78
109+
},
110+
"system_fingerprint": "fp_2a322c9ffc"
111+
}
112+
"
113+
`;
114+
87115
exports[`.parse() zod top-level recursive schemas 1`] = `
88116
"{
89117
"id": "chatcmpl-9uLhw79ArBF4KsQQOlsoE68m6vh6v",

‎tests/lib/parser.test.ts

Lines changed: 175 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -525,13 +525,6 @@ describe('.parse()', () => {
525525
"$schema": "http://json-schema.org/draft-07/schema#",
526526
"additionalProperties": false,
527527
"definitions": {
528-
"contactPerson_properties_person1_properties_name": {
529-
"type": "string",
530-
},
531-
"contactPerson_properties_person1_properties_phone_number": {
532-
"nullable": true,
533-
"type": "string",
534-
},
535528
"query": {
536529
"additionalProperties": false,
537530
"properties": {
@@ -616,6 +609,21 @@ describe('.parse()', () => {
616609
},
617610
],
618611
},
612+
"query_properties_fields_items_anyOf_0_properties_metadata_anyOf_0": {
613+
"additionalProperties": false,
614+
"properties": {
615+
"foo": {
616+
"$ref": "#/definitions/query_properties_fields_items_anyOf_0_properties_metadata_anyOf_0_properties_foo",
617+
},
618+
},
619+
"required": [
620+
"foo",
621+
],
622+
"type": "object",
623+
},
624+
"query_properties_fields_items_anyOf_0_properties_metadata_anyOf_0_properties_foo": {
625+
"type": "string",
626+
},
619627
},
620628
"properties": {
621629
"fields": {
@@ -783,5 +791,165 @@ describe('.parse()', () => {
783791
}
784792
`);
785793
});
794+
795+
test('recursive schema extraction', async () => {
796+
const baseLinkedListNodeSchema = z.object({
797+
value: z.number(),
798+
});
799+
type LinkedListNode = z.infer<typeof baseLinkedListNodeSchema> & {
800+
next: LinkedListNode | null;
801+
};
802+
const linkedListNodeSchema: z.ZodType<LinkedListNode> = baseLinkedListNodeSchema.extend({
803+
next: z.lazy(() => z.union([linkedListNodeSchema, z.null()])),
804+
});
805+
806+
// Define the main schema
807+
const mainSchema = z.object({
808+
linked_list: linkedListNodeSchema,
809+
});
810+
811+
expect(zodResponseFormat(mainSchema, 'query').json_schema.schema).toMatchInlineSnapshot(`
812+
{
813+
"$schema": "http://json-schema.org/draft-07/schema#",
814+
"additionalProperties": false,
815+
"definitions": {
816+
"query": {
817+
"additionalProperties": false,
818+
"properties": {
819+
"linked_list": {
820+
"additionalProperties": false,
821+
"properties": {
822+
"next": {
823+
"anyOf": [
824+
{
825+
"$ref": "#/definitions/query_properties_linked_list",
826+
},
827+
{
828+
"type": "null",
829+
},
830+
],
831+
},
832+
"value": {
833+
"type": "number",
834+
},
835+
},
836+
"required": [
837+
"value",
838+
"next",
839+
],
840+
"type": "object",
841+
},
842+
},
843+
"required": [
844+
"linked_list",
845+
],
846+
"type": "object",
847+
},
848+
"query_properties_linked_list": {
849+
"additionalProperties": false,
850+
"properties": {
851+
"next": {
852+
"$ref": "#/definitions/query_properties_linked_list_properties_next",
853+
},
854+
"value": {
855+
"$ref": "#/definitions/query_properties_linked_list_properties_value",
856+
},
857+
},
858+
"required": [
859+
"value",
860+
"next",
861+
],
862+
"type": "object",
863+
},
864+
"query_properties_linked_list_properties_next": {
865+
"anyOf": [
866+
{
867+
"$ref": "#/definitions/query_properties_linked_list",
868+
},
869+
{
870+
"type": "null",
871+
},
872+
],
873+
},
874+
"query_properties_linked_list_properties_value": {
875+
"type": "number",
876+
},
877+
},
878+
"properties": {
879+
"linked_list": {
880+
"additionalProperties": false,
881+
"properties": {
882+
"next": {
883+
"anyOf": [
884+
{
885+
"$ref": "#/definitions/query_properties_linked_list",
886+
},
887+
{
888+
"type": "null",
889+
},
890+
],
891+
},
892+
"value": {
893+
"type": "number",
894+
},
895+
},
896+
"required": [
897+
"value",
898+
"next",
899+
],
900+
"type": "object",
901+
},
902+
},
903+
"required": [
904+
"linked_list",
905+
],
906+
"type": "object",
907+
}
908+
`);
909+
910+
const completion = await makeSnapshotRequest(
911+
(openai) =>
912+
openai.beta.chat.completions.parse({
913+
model: 'gpt-4o-2024-08-06',
914+
messages: [
915+
{
916+
role: 'system',
917+
content:
918+
"You are a helpful assistant. Generate a data model according to the user's instructions.",
919+
},
920+
{ role: 'user', content: 'create a linklist from 1 to 5' },
921+
],
922+
response_format: zodResponseFormat(mainSchema, 'query'),
923+
}),
924+
2,
925+
);
926+
927+
expect(completion.choices[0]?.message).toMatchInlineSnapshot(`
928+
{
929+
"content": "{"linked_list":{"value":1,"next":{"value":2,"next":{"value":3,"next":{"value":4,"next":{"value":5,"next":null}}}}}}",
930+
"parsed": {
931+
"linked_list": {
932+
"next": {
933+
"next": {
934+
"next": {
935+
"next": {
936+
"next": null,
937+
"value": 5,
938+
},
939+
"value": 4,
940+
},
941+
"value": 3,
942+
},
943+
"value": 2,
944+
},
945+
"value": 1,
946+
},
947+
},
948+
"refusal": null,
949+
"role": "assistant",
950+
"tool_calls": [],
951+
}
952+
`);
953+
});
786954
});
787955
});

0 commit comments

Comments
 (0)