Using Zod with OpenAI Functions

When using the OpenAI GPT API, you need to define functions in a JSON schema, sorta like this:

{
    "name": "get_current_weather",
    "description": "Get the current weather in a given location",
    "parameters": {
        "type": "object",
        "properties": {
            "location": {
                "type": "string",
                "description": "The city and state, e.g. San Francisco, CA",
            },
            "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
        },
        "required": ["location"],
    }
}

If you're using Typescript, this is not great because you manually have to define a type for the parameters of your function and extract them later on.

In my opinion, the ideal workflow would be to define the schema once and "use" it to parse the results you get back from OpenAI. All of this should be type-safe.

We can use Zod for this.

We change how we define functions to use Zod as follows:

const weatherSchema = z.object({
    location: z.string({
        description: "The city and state, e.g. San Francisco, CA",
    }),
    unit: z.enum(["celsius", "fahrenheit"], {
        description: "The unit of temperature to return",
    }).optional()
}, {
    description: "Get the current weather in a given location",
})

const functionDefinitions = {
    get_current_weather: weatherSchema
}

Function definitions is an object, where the keys are the names of the function and the value is a zod schema.

To parse results we get back, we can simply use our function definition:

function handleFunctionCall(functionCall: {
    name: string,
    arguments: string
}) {
    if (functionCall.name === "get_current_weather") {
        const weather = functionDefinitions.get_current_weather.parse(
            JSON.parse(functionCall.arguments)
        )
        
        // weather is now a type-safe object
        // {
        //     location: string
        //     unit?: "celsius" | "fahrenheit"
        // }
    }
}

To turn our function definition into a schema OpenAI understands, I wrote this helper function utilizing the zod-to-json-schema library.

import { ZodSchema } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";

export function functionDefinitionToSchema(
  functions: Record<string, ZodSchema<any>>
) {
  const defintion = [];

  for (const [name, schema] of Object.entries(functions)) {
    const jsonSchema = zodToJsonSchema(schema);
    const description = jsonSchema.description;
    delete jsonSchema.description;

    const functionDefintion = {
      name,
      description,
      parameters: jsonSchema,
    };

    defintion.push(functionDefintion);
  }

  return defintion;
}

I hope that helps you write DRY-er code when utilizing functions in your AI apps. I don't love this approach and I think one could do better so if you have ideas, let me know!