zod-interpreter

Brings support for classes to Zod, the popular TypeScript validation library. This is especially useful for building AI applications, allowing you to write schemas and associated logic all in the same place.

npm install zod-interpreter

One example:

Suppose you would like to make an application that allows users to generate textbooks in collaboration with an AI agent. You might start with some notion of a chapter, which would have four different representations:

1. A UI component rendered on-screen
2. A schema describing the shape of its data to the agent
3. A string representation to be included in the agent's prompt during the editing process
4. A JSON representation to be persisted in the database

A good data structure for this would be a class initialized from JSON whose methods return these different views. And so that's the idea behind z.interpreter(), a new Zod primitive that calls initialize, a special constructor, on its data as part of the parsing operation.

To illustrate, here is a very simple implementation of a Chapter.

import { z } from "zod-interpreter";

const Chapter = z
  .interpreter(
    z.object({
      title: z.string().describe("Succinct title"),
      content: z.string().describe("Content in Markdown"),
    }),
  )
  .initialize(({ _ }) => ({
    _, // <- A property containing the validated data
    render() {
      // Renders some UI component
      return (
        <div>
          <h1>{this._.title}</h1>
          <p>{this._.content}</p>
        </div>
      );
    },
    stringify(chapterNumber) {
      // Returns some string representation for use by agents

      // Here we split by newline and number each chunk so that
      // the agent can direct its edits to particular sections
      return `\
      Chapter ${chapterNumber}: ${this._.title}
      ${this._.content
        .split("\n")
        .map((line, i) => `${i} > ${line}`)
        .join("\n")}`;
    },
  }));

// usage
const chapter = Chapter.parse({
  title: "The Transformer",
  content: "One neural net architecture in particular...",
});

// still works like normal Zod, passing all the same unit tests
Chapter.parse({ bad: "data" }); // throws a ZodError

Chapter implicitly exposes four different representations:

// 1. creates a JSX.Element to display client-side
chapter.render();

// 2. creates a JSON schema to control the agent's outputs
import { zodToJsonSchema } from "zod-to-json-schema";
zodToJsonSchema(Chapter);

// 3. derives a string representation for the agent's prompt
chapter.stringify();

// 4. outputs the current state of the data as JSON
chapter.dump();

Logging chapter prints the following object:

{
  _: {
    title: "The Transformer",
    content: "One neural net architecture in particular..."
  },
  render: function() { ... },
  dump: function() { ... }
}

The validated data is stored in the _ property. You can access it by calling chapter.dump(), returning:

{
  title: "The Transformer",
  content: "One neural net architecture in particular..."
}

Note: you must include the validated data in the object returned by initialize, and it must be contained in the property _. This enables a robust implementation of dump without any further configuration necessary. TypeScript will complain if you forget.

Modularity

Interpreters can be composed, just like any other Zod type.

const Textbook = z
  .interpreter(
    z.object({
      title: z.string(),
      chapters: Chapter.array(),
    }),
  )
  .initialize(({ _ }) => ({
    _,
    render() {
      // More complicated, nested UI
      return (
        <div>
          <h1>{this._.title}</h1>
          {this._.chapters.map((chapter, i) => (
            <div key={i}>{chapter.render()}</div>
          ))}
        </div>
      );
    },
    stringify() {
      return `\
      Textbook title: ${this._.title}
      ${this._.chapters.map((chapter, i) => chapter.stringify(i)).join("\n")}`;
    },
  }));

Useful patterns for AI applications

Encapsulating logic in classes dramatically simplifies the process of defining AI interactions.

For instance, suppose we would like to make the textbook directly editable by the agent. We can accomplish this by exposing a simple editing API:

const Chapter = z
  .interpreter(
    z.object({
      title: z.string().describe("Succinct title"),
      content: z.string().describe("Content in Markdown"),
    }),
  )
  .initialize(({ _ }) => ({
    _,
    render() {
      return (
        <div>
          <h1>{this._.title}</h1>
          <p>{this._.content}</p>
        </div>
      );
    },
    transformContent(transformer) {
      return this._.content.split("\n").map(transformer).join("\n");
    },
    editContent(chunkIndex, newChunk) {
      this._.content = this._.transformContent((chunk, i) => {
        return i === chunkIndex ? newChunk : chunk;
      });
    },
    stringify(chapterIndex) {
      return `\
      Chapter ${chapterIndex}: ${this._.title}
      ${this._.transformContent((chunk, i) => {
        return `${i} > ${chunk}`;
      })}`;
    },
  }));

const Textbook = z
  .interpreter(
    z.object({
      title: z.string().describe("Something catchy"),
      chapters: Chapter.array(),
    }),
  )
  .initialize(({ _ }) => ({
    _,
    render() {
      return (
        <div>
          <h1>{this._.title}</h1>
          {this._.chapters.map((chapter, i) => (
            <div key={i}>{chapter.render()}</div>
          ))}
        </div>
      );
    },
    stringify() {
      return `\
      Title: ${this._.title}
      -----------------------------------------------
      ${this._.chapters
        .map((chapter, i) => chapter.stringify(i))
        .join("\n-----------------------------------------------\n")}`;
    },
    edit(chapterIndex, chunkIndex, newChunk) {
      this._.chapters[chapterIndex].editContent(chunkIndex, newChunk);
    },
  }));

Now our agent can edit the textbook simply by writing an array of objects that contain chapterIndex, chunkIndex, newChunk, much in the same way we might edit a page by rewriting particular paragraphs.

An elegant way to stream these edits to the client as they are written is to render the entire UI as a server component, and swap it out every time an edit is requested with a copy that can be written to directly on the server.

Vercel's ai/rsc library makes this relatively straightforward. We will use it to define a submitUserMessage function that takes a message and returns a stream of UI components.

Our client component that manages the user's requests might look something like this:

// app/components/editor.tsx
"use client";

import { useActions } from "ai/rsc";

export function Editor({ initialUI }: { initialUI: React.ReactNode }) {
  const [textbookUI, setTextbookUI] = useState(initialUI);
  const { submitUserMessage } = useActions();
  return (
    <>
      <div>{textbookUI}</div>
      <Chat
        onSubmit={async (userMessage) => {
          const { ui } = await submitUserMessage(userMessage);

          // Swap the UI with a copy that is being edited by the agent on the server
          setTextbookUI(ui);
        }}
      />
    </>
  );
}

We can render this inside of a Page server component, which will fetch the current state of the textbook from the database as JSON, and then interpret that JSON as initialUI.

// app/page.tsx

export default async function Page() {
  const textbookJson = await getTextbook();
  const textbook = Textbook.parse(textbookJson);
  return <Editor initialUI={textbook.render()} />;
}

Now we just have to define our server-side logic for generating and applying edits to the textbook.

import { openai } from "@ai-sdk/openai";
import { experimental_generateObject } from "ai";
import { getMutableAIState, createStreamableUI } from "ai/rsc";

async function submitUserMessage(content: string) {
  // Boilerplate, see the ai/rsc demo for more details
  const state = getMutableAIState();
  state.update([
    ...state.get(),
    { messages: [...state.get().messages, { role: "user", content }] },
  ]);

  // Create a UI stream, initialized with the current textbook UI
  const textbook = Textbook.parse(state.get().textbookJson);
  const textbookUI = createStreamableUI(textbook.render());

  // Generate edits without blocking the main thread.
  // For simplicity we will only send edits to the client when they're all finished.
  (async () => {
    const { object } = await experimental_generateObject({
      model: openai.chat("gpt-3.5-turbo"),
      schema: z
        .object({
          chapterIndex: z
            .string()
            .describe("The index of the chapter you'd like to edit"),
          chunkIndex: z
            .string()
            .describe("The index of a chunk you'd like to overwrite"),
          newChunk: z.string().describe("The new chunk of text"),
        })
        .array(),
      prompt: `\
      You are a textbook editor. Please make any edits requested in the conversation below.
      Here is the current draft:
      ${textbook.stringify()}
      Here is the conversation history:
      ${state
        .get()
        .messages.map(({ role, content }) => `${role}: ${content}`)
        .join("\n")}
      `,
    });

    // Apply generated edits
    object.forEach(({ chapterIndex, chunkIndex, newChunk }) => {
      textbook.edit(chapterIndex, chunkIndex, newChunk);
    });

    // Close streams with their final values
    textbookUI.done(textbook.render());
    state.done({
      ...state.get(),
      textbookJson: textbook.dump(),
    });

    // Persist result
    await db.insert(textbook.dump());
  })();

  // Immediately return the current UI, before the agent has finished
  return {
    id: new Date().toISOString(),
    ui: textbookUI.value,
  };
}

Explicit typing

TypeScript will infer the correct types of the outputs of each method, but not the inputs.

To explicitly type the interpreter output, you can provide type arguments to the initialize method, which takes a mapped type over each derived property. You do not need to type the validated data stored in _.

If TypeScript ever complains about "excessive depth" in its type inference, the solution is just to manually provide types for any particularly complex interpreters.

const Textbook = z
  .interpreter(
    z.object({
      title: z.string().describe("Something catchy"),
      chapters: Chapter.array(),
    }),
  )
  .initialize<{
    render: () => React.ReactNode; 
    stringify: () => string; 
    edit: (chapterIndex: number, chunkIndex: number, newChunk: string) => void; 
  }>(({ _ }) => ({
    _,
    render() {
      return (
        <div>
          <h1>{this._.title}</h1>
          {this._.chapters.map((chapter, i) => (
            <div key={i}>{chapter.render()}</div>
          ))}
        </div>
      );
    },
    stringify() {
      return `\
      Title: ${this._.title}
      -----------------------------------------------
      ${this._.chapters
        .map((chapter, i) => chapter.stringify(i))
        .join("\n-----------------------------------------------\n")}`;
    },
    edit(chapterIndex, chunkIndex, newChunk) {
      this._.chapters[chapterIndex].editContent(chunkIndex, newChunk);
    },
  }));