fix(helpers/zod): add extract-to-root ref strategy

openai · Aug 8, 2024 · ef3c73c · ef3c73c
1 parent e4a247a
commit ef3c73c
Show file tree

Hide file tree

Showing 7 changed files with 69 additions and 28 deletions.
diff --git a/src/_vendor/zod-to-json-schema/Options.ts b/src/_vendor/zod-to-json-schema/Options.ts
@@ -10,7 +10,7 @@ export const ignoreOverride = Symbol('Let zodToJsonSchema decide on which parser
 
 export type Options<Target extends Targets = 'jsonSchema7'> = {
   name: string | undefined;
-  $refStrategy: 'root' | 'relative' | 'none' | 'seen';
+  $refStrategy: 'root' | 'relative' | 'none' | 'seen' | 'extract-to-root';
   basePath: string[];
   effectStrategy: 'input' | 'any';
   pipeStrategy: 'input' | 'output' | 'all';
@@ -20,7 +20,7 @@ export type Options<Target extends Targets = 'jsonSchema7'> = {
   target: Target;
   strictUnions: boolean;
   definitionPath: string;
-  definitions: Record<string, ZodSchema>;
+  definitions: Record<string, ZodSchema | ZodTypeDef>;
   errorMessages: boolean;
   markdownDescription: boolean;
   patternStrategy: 'escape' | 'preserve';

diff --git a/src/_vendor/zod-to-json-schema/Refs.ts b/src/_vendor/zod-to-json-schema/Refs.ts
@@ -1,6 +1,7 @@
-import { ZodTypeDef } from 'zod';
+import type { ZodTypeDef } from 'zod';
 import { getDefaultOptions, Options, Targets } from './Options';
 import { JsonSchema7Type } from './parseDef';
+import { zodDef } from './util';
 
 export type Refs = {
   seen: Map<ZodTypeDef, Seen>;
@@ -33,9 +34,9 @@ export const getRefs = (options?: string | Partial<Options<Targets>>): Refs => {
     seenRefs: new Set(),
     seen: new Map(
       Object.entries(_options.definitions).map(([name, def]) => [
-        def._def,
+        zodDef(def),
         {
-          def: def._def,
+          def: zodDef(def),
           path: [..._options.basePath, _options.definitionPath, name],
           // Resolution of references will be forced even though seen, so it's ok that the schema is undefined here for now.
           jsonSchema: undefined,

diff --git a/src/_vendor/zod-to-json-schema/parseDef.ts b/src/_vendor/zod-to-json-schema/parseDef.ts
@@ -122,6 +122,24 @@ const get$ref = (
   switch (refs.$refStrategy) {
     case 'root':
       return { $ref: item.path.join('/') };
+    // this case is needed as OpenAI strict mode doesn't support top-level `$ref`s, i.e.
+    // the top-level schema *must* be `{"type": "object", "properties": {...}}` but if we ever
+    // need to define a `$ref`, relative `$ref`s aren't supported, so we need to extract
+    // the schema to `#/definitions/` and reference that.
+    //
+    // e.g. if we need to reference a schema at
+    // `["#","definitions","contactPerson","properties","person1","properties","name"]`
+    // then we'll extract it out to `contactPerson_properties_person1_properties_name`
+    case 'extract-to-root':
+      const name = item.path.slice(refs.basePath.length + 1).join('_');
+
+      // we don't need to extract the root schema in this case, as it's already
+      // been added to the definitions
+      if (name !== refs.name && refs.nameStrategy === 'duplicate-ref') {
+        refs.definitions[name] = item.def;
+      }
+
+      return { $ref: [...refs.basePath, refs.definitionPath, name].join('/') };
     case 'relative':
       return { $ref: getRelativePath(refs.currentPath, item.path) };
     case 'none':

diff --git a/src/_vendor/zod-to-json-schema/util.ts b/src/_vendor/zod-to-json-schema/util.ts
@@ -0,0 +1,11 @@
+import type { ZodSchema, ZodTypeDef } from 'zod';
+
+export const zodDef = (zodSchema: ZodSchema | ZodTypeDef): ZodTypeDef => {
+  return '_def' in zodSchema ? zodSchema._def : zodSchema;
+};
+
+export function isEmptyObj(obj: Object | null | undefined): boolean {
+  if (!obj) return true;
+  for (const _k in obj) return false;
+  return true;
+}
diff --git a/src/_vendor/zod-to-json-schema/zodToJsonSchema.ts b/src/_vendor/zod-to-json-schema/zodToJsonSchema.ts
@@ -2,6 +2,7 @@ import { ZodSchema } from 'zod';
 import { Options, Targets } from './Options';
 import { JsonSchema7Type, parseDef } from './parseDef';
 import { getRefs } from './Refs';
+import { zodDef, isEmptyObj } from './util';
 
 const zodToJsonSchema = <Target extends Targets = 'jsonSchema7'>(
   schema: ZodSchema<any>,
@@ -16,25 +17,6 @@ const zodToJsonSchema = <Target extends Targets = 'jsonSchema7'>(
 } => {
   const refs = getRefs(options);
 
-  const definitions =
-    typeof options === 'object' && options.definitions ?
-      Object.entries(options.definitions).reduce(
-        (acc, [name, schema]) => ({
-          ...acc,
-          [name]:
-            parseDef(
-              schema._def,
-              {
-                ...refs,
-                currentPath: [...refs.basePath, refs.definitionPath, name],
-              },
-              true,
-            ) ?? {},
-        }),
-        {},
-      )
-    : undefined;
-
   const name =
     typeof options === 'string' ? options
     : options?.nameStrategy === 'title' ? undefined
@@ -61,6 +43,25 @@ const zodToJsonSchema = <Target extends Targets = 'jsonSchema7'>(
     main.title = title;
   }
 
+  const definitions =
+    !isEmptyObj(refs.definitions) ?
+      Object.entries(refs.definitions).reduce(
+        (acc, [name, schema]) => ({
+          ...acc,
+          [name]:
+            parseDef(
+              zodDef(schema),
+              {
+                ...refs,
+                currentPath: [...refs.basePath, refs.definitionPath, name],
+              },
+              true,
+            ) ?? {},
+        }),
+        {},
+      )
+    : undefined;
+
   const combined: ReturnType<typeof zodToJsonSchema<Target>> =
     name === undefined ?
       definitions ?

diff --git a/src/helpers/zod.ts b/src/helpers/zod.ts
@@ -13,6 +13,7 @@ function zodToJsonSchema(schema: z.ZodType, options: { name: string }): Record<s
     openaiStrictMode: true,
     name: options.name,
     nameStrategy: 'duplicate-ref',
+    $refStrategy: 'extract-to-root',
   });
 }
 

diff --git a/tests/lib/parser.test.ts b/tests/lib/parser.test.ts
@@ -352,10 +352,10 @@ describe('.parse()', () => {
                       "type": "string",
                     },
                     "name": {
-                      "$ref": "#/definitions/contactPerson/properties/person1/properties/name",
+                      "$ref": "#/definitions/contactPerson_properties_person1_properties_name",
                     },
                     "phone_number": {
-                      "$ref": "#/definitions/contactPerson/properties/person1/properties/phone_number",
+                      "$ref": "#/definitions/contactPerson_properties_person1_properties_phone_number",
                     },
                   },
                   "required": [
@@ -372,6 +372,15 @@ describe('.parse()', () => {
               ],
               "type": "object",
             },
+            "contactPerson_properties_person1_properties_name": {
+              "type": "string",
+            },
+            "contactPerson_properties_person1_properties_phone_number": {
+              "type": [
+                "string",
+                "null",
+              ],
+            },
           },
           "properties": {
             "person1": {
@@ -424,10 +433,10 @@ describe('.parse()', () => {
                   "type": "string",
                 },
                 "name": {
-                  "$ref": "#/definitions/contactPerson/properties/person1/properties/name",
+                  "$ref": "#/definitions/contactPerson_properties_person1_properties_name",
                 },
                 "phone_number": {
-                  "$ref": "#/definitions/contactPerson/properties/person1/properties/phone_number",
+                  "$ref": "#/definitions/contactPerson_properties_person1_properties_phone_number",
                 },
               },
               "required": [