fix: automatic zod schema creation for ai sdk provider tools

Author: dosco

package-lock.json

@@ -19,10 +19,12 @@ import type {
   AxChatResponse,
   AxChatResponseResult,
   AxFunction,
-  AxFunctionJSONSchema,
   AxGenIn,
   AxGenOut
 } from '@ax-llm/ax/index.js';
+import { z } from 'zod';
+
+import { convertToZodSchema } from './util.js';
 
 type Writeable<T> = { -readonly [P in keyof T]: T[P] };
 type AxChatRequestChatPrompt = Writeable<AxChatRequest['chatPrompt'][0]>;
@@ -33,15 +35,9 @@ type AxConfig = {
 
 type generateFunction<T> = ((input: T) => Promise<unknown>) | undefined;
 
-const schemaSymbol = Symbol('vercel.ai.schema');
-
 interface RenderTool<IN> {
   description?: string;
-  parameters?: {
-    [schemaSymbol]: true;
-    validate: undefined;
-    jsonSchema: AxFunctionJSONSchema;
-  };
+  parameters?: z.ZodTypeAny;
   generate?: generateFunction<IN>;
 }
 
@@ -72,11 +68,7 @@ export class AxAgentProvider<IN extends AxGenIn, OUT extends AxGenOut>
       properties: {}
     };
 
-    return {
-      [schemaSymbol]: true as const,
-      validate: undefined,
-      jsonSchema
-    };
+    return convertToZodSchema(jsonSchema);
   }
 
   get generate(): generateFunction<IN> {

src/ai-sdk-provider/index.ts

@@ -24,7 +24,8 @@
   },
   "dependencies": {
     "@ai-sdk/provider-utils": "^1.0.2",
-    "@ax-llm/ax": "9.0.39"
+    "@ax-llm/ax": "9.0.39",
+    "zod": "^3.23.8"
   },
   "devDependencies": {
     "npm-run-all": "^4.1.5",

src/ai-sdk-provider/package.json

@@ -0,0 +1,44 @@
+import type { AxFunctionJSONSchema } from '@ax-llm/ax';
+import { z } from 'zod';
+
+type AnyZod =
+  | z.AnyZodObject
+  | z.ZodString
+  | z.ZodNumber
+  | z.ZodBoolean
+  | z.ZodArray<AnyZod>
+  | z.ZodOptional<AnyZod>;
+
+export function convertToZodSchema(
+  jsonSchema: Readonly<AxFunctionJSONSchema>
+): AnyZod {
+  const { type, properties, required, items } = jsonSchema;
+
+  switch (type) {
+    case 'string':
+      return z.string();
+    case 'number':
+      return z.number();
+    case 'boolean':
+      return z.boolean();
+    case 'array':
+      if (!items) {
+        throw new Error("Array type must have 'items' property.");
+      }
+      return z.array(convertToZodSchema(items));
+    case 'object': {
+      if (!properties) {
+        throw new Error("Object type must have 'properties' property.");
+      }
+      const shape: Record<string, AnyZod> = {};
+
+      for (const [key, value] of Object.entries(properties)) {
+        const schema = convertToZodSchema(value);
+        shape[key] = required?.includes(key) ? schema : schema.optional();
+      }
+      return z.object(shape);
+    }
+    default:
+      throw new Error(`Unsupported type: ${type}`);
+  }
+}