AI Assistant Output with End to End Type Safety

As I wrote about previously, I’ve been using structured output from AI Assistants to generate app features. One issue that emerged right away was the need to manually create Typescript types matching the assistant output instructions, as can be exemplified from the following:

- `questions`: An array of trivia questions in a mix of common trivia categories. Each question
   should be unique within the array and contain an object with the following properties.
  - `category`: The trivia category.
  - `question`: The trivia question.
  - `choices`: The answer choices.
  - `correctAnswer`: The correct answer.

export interface TriviaGame {
  questions: TriviaQuestion[];
}

export interface TriviaQuestion {
  question: string;
  category: string;
  choices: string[];
  correctAnswer: string;
}

Both of the above are describing the identical data structure but in two different ways. Obviously, repeating the same information in multiple places is something we want to avoid.

In this post, I’ll talk about how I replaced repeated instances of the same data definition with a single schema, which then can be used for output instructions, Typescript types, and—as an added bonus—can also be used for runtime type validation.

These are the main steps I took for adding end-to-end typing:

1. Create a schema to support the output requirements
2. Use .describe() as mini assistant instructions
3. Convert to OpenAI Response Format
4. Convert to Typescript
5. Add runtime type validation

Let’s walk through each in detail.

Create a schema to support the output requirements

The feature we’re building is a Trivia Game and the requirements for the assistant output is to generate a set of trivia questions. We’ll be using Zod for defining our schema.

Here are the above JSON Output instructions converted into a Zod schema:

import { z } from "zod";

export const triviaGameSchema = z.object({
  questions: z
    .object({
      category: z.string().describe("The question category."),
      question: z.string().describe("The trivia question."),
      choices: z
        .string()
        .array()
        .describe(
          "The answer choices, which should include one correct answer and three distractor choices.",
        ),
      correctAnswer: z.string().describe("The correct answer."),
    })
    .array()
    .describe(
      "An array of trivia questions in a mix of common trivia categories. Each question should be unique within the array.",
    ),
});

If you’ve used Zod before, then the above should look quite familiar. However, if you’re new to creating schemas and/or using Zod, I suggest reviewing their docs if any of the above is unclear. It’s worth the effort, in particular, since, as we’ll see, we’re going to get quite a bit of mileage out of this one schema.

Use .describe() as mini assistant instructions

One part of the schema which may look unfamiliar, even if you’ve used Zod previously, is the use of the describe method. At least for me, it wasn’t something I’d had a need for until I was using it for structured output.

However, in the context of structured output, this property is essential. One can think of these statements as attribute-level assistant instructions. Use this to explain requirements for an attribute, particularly those that are not self-evident from the attribute name or the data structure itself. One example from above is where we state that the question array should contain questions in “a mix of common trivia categories”.

Convert to OpenAI Response Format

With our schema in place, we now want to reuse it in our assistant instructions. To achive this, we can make use of OpenAI’s zodResponseFormat helper.

With that, we can then use the schema as follows to create (or update) an assistant:

import { zodResponseFormat } from "openai/helpers/zod";

...

const assistant = await openai.beta.assistants.create({
  model: "gpt-4o-mini", // be sure to use a model that supports structured output
  ...
  response_format: zodResponseFormat(mySchema, "my_schema"),
});

As can be seen from above, reusing the zod schema for OpenAI assistant instructions is quite straightforward. However, there are some limitations and issues to be aware of when using this helper.

Convert the schema to Typescript

We also want to reuse our schema for our app types, e.g. for typing props used in components that render the assistant output. This can be achieved as follows:

const triviaGameSchema = z.object({
16 collapsed lines
  questions: z
    .object({
      category: z.string().describe("The question category."),
      question: z.string().describe("The trivia question."),
      choices: z
        .string()
        .array()
        .describe(
          "The answer choices, which should include one correct answer and three distractor choices.",
        ),
      correctAnswer: z.string().describe("The correct answer."),
    })
    .array()
    .describe(
      "An array of trivia questions. Each question should be unique within the array. Question categories should be a mix of common trivia categories.",
    ),
});

type TriviaGame = z.infer<typeof triviaGameSchema>;

With the above, we’ve now eliminated the duplication I discussed in the intro. 🎉

Add runtime type validation

While adding Typescript types is great for development, we still need to also ensure our output is typed as expected during runtime. (In case you weren’t aware, most Typescript types don’t exist at runtime.)

Because we already did the heavy lifing when creating the schema, we can leverage Zod, and validate assistant output using the parse method.

    const output = await getAsstOutput({
   ...
    });

    const game = triviaGameSchema.parse(output);

This will ensure that output returned by the assistant matches what we expect. To be clear, it’s unlikely you’ll get a validation error here, since the assistant is using the same schema as a basis for its response. However, mainly because OpenAI is a third party service outside of our control, meaning we can’t be certain what the response will be, I think it’s good practice to validate the response.

Updating and Creating Assistants all in one place

You might have noticed from the above code that we’re creating our assistant as part of the app code, rather than via the OpenIA dashboard. This is a requirement if we want to be able to reuse our schema both for our app types as well as the instruction responses.

However, creating assistants using using the OpenAI API rather than the dashboard is not as straightforward as one might expect. In the next post, I’ll walk through some challenges you might encounter, as well as my approach for enabling code-based assistant management.