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>
openai · Aug 13, 2024 · ac309ab · ac309ab
1 parent d486d27
commit ac309ab
Show file tree

Hide file tree

Showing 4 changed files with 242 additions and 27 deletions.
diff --git a/src/_vendor/zod-to-json-schema/Options.ts b/src/_vendor/zod-to-json-schema/Options.ts
@@ -38,10 +38,9 @@ export type Options<Target extends Targets = 'jsonSchema7'> = {
  openaiStrictMode?: boolean;
 };
 
-export const defaultOptions: Options = {
+const defaultOptions: Omit<Options, 'definitions' | 'basePath'> = {
  name: undefined,
  $refStrategy: 'root',
- basePath: ['#'],
  effectStrategy: 'input',
  pipeStrategy: 'all',
  dateStrategy: 'format:date-time',
@@ -51,7 +50,6 @@ export const defaultOptions: Options = {
  definitionPath: 'definitions',
  target: 'jsonSchema7',
  strictUnions: false,
- definitions: {},
  errorMessages: false,
  markdownDescription: false,
  patternStrategy: 'escape',
@@ -63,13 +61,20 @@ export const defaultOptions: Options = {
 
 export const getDefaultOptions = <Target extends Targets>(
  options: Partial<Options<Target>> | string | undefined,
-) =>
- (typeof options === 'string' ?
- {
- ...defaultOptions,
- name: options,
- }
- : {
- ...defaultOptions,
- ...options,
- }) as Options<Target>;
+) => {
+ // We need to add `definitions` here as we may mutate it
+ return (
+ typeof options === 'string' ?
+ {
+ ...defaultOptions,
+ basePath: ['#'],
+ definitions: {},
+ name: options,
+ }
+ : {
+ ...defaultOptions,
+ basePath: ['#'],
+ definitions: {},
+ ...options,
+ }) as Options<Target>;
+};
diff --git a/src/_vendor/zod-to-json-schema/zodToJsonSchema.ts b/src/_vendor/zod-to-json-schema/zodToJsonSchema.ts
@@ -49,14 +49,28 @@ const zodToJsonSchema = <Target extends Targets = 'jsonSchema7'>(
  }
 
  const definitions: Record<string, any> = {};
+ const processedDefinitions = new Set();
 
- for (const [name, zodSchema] of Object.entries(refs.definitions)) {
- definitions[name] =
- parseDef(
- zodDef(zodSchema),
- { ...refs, currentPath: [...refs.basePath, refs.definitionPath, name] },
- true,
- ) ?? {};
+ // the call to `parseDef()` here might itself add more entries to `.definitions`
+ // so we need to continually evaluate definitions until we've resolved all of them
+ //
+ // we have a generous iteration limit here to avoid blowing up the stack if there
+ // are any bugs that would otherwise result in us iterating indefinitely
+ for (let i = 0; i < 500; i++) {
+ const newDefinitions = Object.entries(refs.definitions).filter(
+ ([key]) => !processedDefinitions.has(key),
+ );
+ if (newDefinitions.length === 0) break;
+
+ for (const [key, schema] of newDefinitions) {
+ definitions[key] =
+ parseDef(
+ zodDef(schema),
+ { ...refs, currentPath: [...refs.basePath, refs.definitionPath, key] },
+ true,
+ ) ?? {};
+ processedDefinitions.add(key);
+ }
  }
 
  return definitions;

diff --git a/tests/lib/__snapshots__/parser.test.ts.snap b/tests/lib/__snapshots__/parser.test.ts.snap
@@ -84,6 +84,34 @@ exports[`.parse() zod nested schema extraction 2`] = `
 "
 `;
 
+exports[`.parse() zod recursive schema extraction 2`] = `
+"{
+ "id": "chatcmpl-9vdbw9dekyUSEsSKVQDhTxA2RCxcK",
+ "object": "chat.completion",
+ "created": 1723523988,
+ "model": "gpt-4o-2024-08-06",
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "{\\"linked_list\\":{\\"value\\":1,\\"next\\":{\\"value\\":2,\\"next\\":{\\"value\\":3,\\"next\\":{\\"value\\":4,\\"next\\":{\\"value\\":5,\\"next\\":null}}}}}}",
+ "refusal": null
+ },
+ "logprobs": null,
+ "finish_reason": "stop"
+ }
+ ],
+ "usage": {
+ "prompt_tokens": 40,
+ "completion_tokens": 38,
+ "total_tokens": 78
+ },
+ "system_fingerprint": "fp_2a322c9ffc"
+}
+"
+`;
+
 exports[`.parse() zod top-level recursive schemas 1`] = `
 "{
  "id": "chatcmpl-9uLhw79ArBF4KsQQOlsoE68m6vh6v",

diff --git a/tests/lib/parser.test.ts b/tests/lib/parser.test.ts
@@ -525,13 +525,6 @@ describe('.parse()', () => {
  "$schema": "http://json-schema.org/draft-07/schema#",
  "additionalProperties": false,
  "definitions": {
- "contactPerson_properties_person1_properties_name": {
- "type": "string",
- },
- "contactPerson_properties_person1_properties_phone_number": {
- "nullable": true,
- "type": "string",
- },
  "query": {
  "additionalProperties": false,
  "properties": {
@@ -616,6 +609,21 @@ describe('.parse()', () => {
  },
  ],
  },
+ "query_properties_fields_items_anyOf_0_properties_metadata_anyOf_0": {
+ "additionalProperties": false,
+ "properties": {
+ "foo": {
+ "$ref": "#/definitions/query_properties_fields_items_anyOf_0_properties_metadata_anyOf_0_properties_foo",
+ },
+ },
+ "required": [
+ "foo",
+ ],
+ "type": "object",
+ },
+ "query_properties_fields_items_anyOf_0_properties_metadata_anyOf_0_properties_foo": {
+ "type": "string",
+ },
  },
  "properties": {
  "fields": {
@@ -783,5 +791,165 @@ describe('.parse()', () => {
  }
  `);
  });
+
+ test('recursive schema extraction', async () => {
+ const baseLinkedListNodeSchema = z.object({
+ value: z.number(),
+ });
+ type LinkedListNode = z.infer<typeof baseLinkedListNodeSchema> & {
+ next: LinkedListNode | null;
+ };
+ const linkedListNodeSchema: z.ZodType<LinkedListNode> = baseLinkedListNodeSchema.extend({
+ next: z.lazy(() => z.union([linkedListNodeSchema, z.null()])),
+ });
+
+ // Define the main schema
+ const mainSchema = z.object({
+ linked_list: linkedListNodeSchema,
+ });
+
+ expect(zodResponseFormat(mainSchema, 'query').json_schema.schema).toMatchInlineSnapshot(`
+ {
+ "$schema": "http://json-schema.org/draft-07/schema#",
+ "additionalProperties": false,
+ "definitions": {
+ "query": {
+ "additionalProperties": false,
+ "properties": {
+ "linked_list": {
+ "additionalProperties": false,
+ "properties": {
+ "next": {
+ "anyOf": [
+ {
+ "$ref": "#/definitions/query_properties_linked_list",
+ },
+ {
+ "type": "null",
+ },
+ ],
+ },
+ "value": {
+ "type": "number",
+ },
+ },
+ "required": [
+ "value",
+ "next",
+ ],
+ "type": "object",
+ },
+ },
+ "required": [
+ "linked_list",
+ ],
+ "type": "object",
+ },
+ "query_properties_linked_list": {
+ "additionalProperties": false,
+ "properties": {
+ "next": {
+ "$ref": "#/definitions/query_properties_linked_list_properties_next",
+ },
+ "value": {
+ "$ref": "#/definitions/query_properties_linked_list_properties_value",
+ },
+ },
+ "required": [
+ "value",
+ "next",
+ ],
+ "type": "object",
+ },
+ "query_properties_linked_list_properties_next": {
+ "anyOf": [
+ {
+ "$ref": "#/definitions/query_properties_linked_list",
+ },
+ {
+ "type": "null",
+ },
+ ],
+ },
+ "query_properties_linked_list_properties_value": {
+ "type": "number",
+ },
+ },
+ "properties": {
+ "linked_list": {
+ "additionalProperties": false,
+ "properties": {
+ "next": {
+ "anyOf": [
+ {
+ "$ref": "#/definitions/query_properties_linked_list",
+ },
+ {
+ "type": "null",
+ },
+ ],
+ },
+ "value": {
+ "type": "number",
+ },
+ },
+ "required": [
+ "value",
+ "next",
+ ],
+ "type": "object",
+ },
+ },
+ "required": [
+ "linked_list",
+ ],
+ "type": "object",
+ }
+ `);
+
+ const completion = await makeSnapshotRequest(
+ (openai) =>
+ openai.beta.chat.completions.parse({
+ model: 'gpt-4o-2024-08-06',
+ messages: [
+ {
+ role: 'system',
+ content:
+ "You are a helpful assistant. Generate a data model according to the user's instructions.",
+ },
+ { role: 'user', content: 'create a linklist from 1 to 5' },
+ ],
+ response_format: zodResponseFormat(mainSchema, 'query'),
+ }),
+ 2,
+ );
+
+ expect(completion.choices[0]?.message).toMatchInlineSnapshot(`
+ {
+ "content": "{"linked_list":{"value":1,"next":{"value":2,"next":{"value":3,"next":{"value":4,"next":{"value":5,"next":null}}}}}}",
+ "parsed": {
+ "linked_list": {
+ "next": {
+ "next": {
+ "next": {
+ "next": {
+ "next": null,
+ "value": 5,
+ },
+ "value": 4,
+ },
+ "value": 3,
+ },
+ "value": 2,
+ },
+ "value": 1,
+ },
+ },
+ "refusal": null,
+ "role": "assistant",
+ "tool_calls": [],
+ }
+ `);
+ });
  });
 });