feat: ✨ add support for Zod types in schema utilities

microsoft · Dec 21, 2024 · 901527a · 901527a
1 parent 7317c04
commit 901527a
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 17 deletions.
diff --git a/packages/core/src/promptdom.ts b/packages/core/src/promptdom.ts
@@ -362,13 +362,13 @@ export function createImageNode(
 // Function to create a schema node.
 export function createSchemaNode(
     name: string,
-    value: JSONSchema,
+    value: JSONSchema | ZodTypeLike,
     options?: DefSchemaOptions
 ): PromptSchemaNode {
     assert(!!name)
     assert(value !== undefined)
     // auto zod conversion
-    value = tryZodToJsonSchema(value) ?? value
+    value = tryZodToJsonSchema(value as ZodTypeLike) ?? (value as JSONSchema)
     return { type: "schema", name, value, options }
 }
 

diff --git a/packages/core/src/types/prompt_template.d.ts b/packages/core/src/types/prompt_template.d.ts
@@ -2604,10 +2604,12 @@ interface McpServerConfig {
 
 type McpServersConfig = Record<string, Omit<McpServerConfig, "id" | "options">>
 
+type ZodTypeLike = { _def: any, safeParse: any, refine: any }
+
 interface ChatGenerationContext extends ChatTurnGenerationContext {
     defSchema(
         name: string,
-        schema: JSONSchema | object,
+        schema: JSONSchema | ZodTypeLike,
         options?: DefSchemaOptions
     ): string
     defImages(

diff --git a/packages/core/src/types/prompt_type.d.ts b/packages/core/src/types/prompt_type.d.ts
@@ -239,7 +239,7 @@ declare function fetchText(
  */
 declare function defSchema(
     name: string,
-    schema: JSONSchema | object,
+    schema: JSONSchema | ZodTypeLike,
     options?: DefSchemaOptions
 ): string
 

diff --git a/packages/core/src/zod.ts b/packages/core/src/zod.ts
@@ -6,17 +6,14 @@ import { zodToJsonSchema as _zodToJsonSchema } from "zod-to-json-schema"
  * @param options
  * @returns
  */
-export function tryZodToJsonSchema(z: object, options?: object): JSONSchema {
-    if (!z) return undefined
-    // instanceof not working, test for some existing methoid
-    if (!(z as any)._def) return undefined
-    try {
-        const schema = _zodToJsonSchema(z as any, {
-            target: "openAi",
-            ...(options || {}),
-        })
-        return structuredClone(schema) as JSONSchema
-    } catch (e) {
-        return undefined
-    }
+export function tryZodToJsonSchema(
+    z: ZodTypeLike,
+    options?: object
+): JSONSchema {
+    if (!z || !z._def || !z.refine || !z.safeParse) return undefined
+    const schema = _zodToJsonSchema(z as any, {
+        target: "openAi",
+        ...(options || {}),
+    })
+    return structuredClone(schema) as JSONSchema
 }