From 901527aa22cf24b667c13a005cf55da8dcd0db10 Mon Sep 17 00:00:00 2001
From: Peli de Halleux <pelikhan@users.noreply.github.com>
Date: Sat, 21 Dec 2024 15:17:49 +0000
Subject: [PATCH] =?UTF-8?q?feat:=20=E2=9C=A8=20add=20support=20for=20Zod?=
 =?UTF-8?q?=20types=20in=20schema=20utilities?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 packages/core/src/promptdom.ts               |  4 ++--
 packages/core/src/types/prompt_template.d.ts |  4 +++-
 packages/core/src/types/prompt_type.d.ts     |  2 +-
 packages/core/src/zod.ts                     | 23 +++++++++-----------
 4 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/packages/core/src/promptdom.ts b/packages/core/src/promptdom.ts
index 0456148a8..71d9f5f63 100644
--- 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
index 6ff01de74..f6481a577 100644
--- 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
index 529770149..56c7042f3 100644
--- 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
index 68f2a94da..f3fe5a600 100644
--- 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
 }