How to set $refStrategy for zodToJsonSchema to use strict OpenAI GPT JSON response mode with typed results

[sibbl]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import { ChatOpenAI } from "@langchain/openai";
import { z } from "zod";
import {
  ChatPromptTemplate,
  HumanMessagePromptTemplate
} from "@langchain/core/prompts";
import { SystemMessage } from "@langchain/core/messages";

const promptTemplate = ChatPromptTemplate.fromMessages([
  new SystemMessage("Be friendly to the user"),
  HumanMessagePromptTemplate.fromTemplate("{userInput}")
]);

type SchemaBaseType = {
  type: "one" | "two";
  label: string;
};

type SchemaOneType = SchemaBaseType & {
  type: "one";
  one: string | null;
};

type SchemaTwoType = SchemaBaseType & {
  type: "two";
  two: string | null;
};

type SchemaType = SchemaOneType | SchemaTwoType;

type ResultType = {
  elements: SchemaType[];
};

async function getCompletionAsync(): Promise<ResultType> {
  const SchemaBase = z
    .object({
      type: z.enum(["one", "two"]),
      label: z.string()
    })
    .strict();

  const SchemaOne = SchemaBase.extend({
    type: z.literal("one"),
    one: z.union([z.string(), z.null()])
  });

  const SchemaTwo = SchemaBase.extend({
    type: z.literal("two"),
    two: z.union([z.string(), z.null()])
  });

  const schema = z.object({
    elements: z.array(z.discriminatedUnion("type", [SchemaOne, SchemaTwo]))
  });

  const model = new ChatOpenAI();
  const structuredLlm = await model.withStructuredOutput(schema, {
    strict: true
  });

  const chain = promptTemplate.pipe(structuredLlm);
  const result = await chain.invoke({
    userInput: "hello world"
  });
  return result;
}

Description

• I'm trying to use the strict mode with a JSON more complex schema in Typescript
• zod internally uses the JSON schema feature of refs to generate a pretty JSON schema
• OpenAI responds with 400 Invalid schema for function 'extract': In context=('properties', 'elements', 'items', 'anyOf', '1', 'properties', 'label'), reference can only point to definitions defined at the top level of the schema.
• internally, zodToJsonSchema is used, which accepts a $refStrategy: "none" option to avoid refs being used
• I'm looking for a way to pass this option to zodToJsonSchema, so that a JSON schema without refs is sent to OpenAI

By using

import { zodToJsonSchema } from "zod-to-json-schema";
...
const jsonSchema = zodToJsonSchema(schema, {
  $refStrategy: "none"
});
const structuredLlm = await model.withStructuredOutput(jsonSchema, {
  strict: true
});

I can basically get it working, but I loose the typed result as model.withStructuredOutput then returns Record<string, any> and not ResultType anymore.

System Info

Windows, Node 20.16.0, npm 10.8.1

"@langchain/community": "0.2.31",
"@langchain/langgraph": "0.0.34",
"@langchain/openai": "0.2.8",

[dosubot[bot]]

Hey, @sibbl! I'm here to help you out with your coding questions and bug fixes. Let's tackle this issue together!

To set the $refStrategy for zodToJsonSchema to use strict OpenAI GPT JSON response mode with typed results, you can follow these steps:

1. Convert the Zod schema to JSON schema with $refStrategy: "none": This avoids the use of references in the JSON schema, which OpenAI's API does not support in certain contexts.
2. Use the JSON schema with withStructuredOutput: This ensures that the model's output matches the schema without losing the type information.

Here's how you can modify your code to achieve this:

import { ChatOpenAI } from "@langchain/openai";
import { z } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";
import {
  ChatPromptTemplate,
  HumanMessagePromptTemplate
} from "@langchain/core/prompts";
import { SystemMessage } from "@langchain/core/messages";

const promptTemplate = ChatPromptTemplate.fromMessages([
  new SystemMessage("Be friendly to the user"),
  HumanMessagePromptTemplate.fromTemplate("{userInput}")
]);

type SchemaBaseType = {
  type: "one" | "two";
  label: string;
};

type SchemaOneType = SchemaBaseType & {
  type: "one";
  one: string | null;
};

type SchemaTwoType = SchemaBaseType & {
  type: "two";
  two: string | null;
};

type SchemaType = SchemaOneType | SchemaTwoType;

type ResultType = {
  elements: SchemaType[];
};

async function getCompletionAsync(): Promise<ResultType> {
  const SchemaBase = z
    .object({
      type: z.enum(["one", "two"]),
      label: z.string()
    })
    .strict();

  const SchemaOne = SchemaBase.extend({
    type: z.literal("one"),
    one: z.union([z.string(), z.null()])
  });

  const SchemaTwo = SchemaBase.extend({
    type: z.literal("two"),
    two: z.union([z.string(), z.null()])
  });

  const schema = z.object({
    elements: z.array(z.discriminatedUnion("type", [SchemaOne, SchemaTwo]))
  });

  // Convert Zod schema to JSON schema with $refStrategy: "none"
  const jsonSchema = zodToJsonSchema(schema, {
    $refStrategy: "none"
  });

  const model = new ChatOpenAI();
  const structuredLlm = await model.withStructuredOutput(jsonSchema, {
    strict: true
  });

  const chain = promptTemplate.pipe(structuredLlm);
  const result = await chain.invoke({
    userInput: "hello world"
  });
  return result as ResultType; // Cast the result to ResultType
}

By converting the Zod schema to JSON schema with $refStrategy: "none", you avoid the use of references, which should resolve the 400 Invalid schema error. Casting the result to ResultType ensures that you retain the type information.

This approach allows you to use strict mode with a complex JSON schema while maintaining typed results.

To continue talking to Dosu, mention @dosu.

[sibbl]

As described above, the result will be of type Record<string, any> and not ResultType. I don't want to cast. There should be a way to use it with proper Typescript support.

If there's not such an option yet, I'd suggest to create a PR to introduce a way to pass options to zodToJsonSchema or automatically set the $refStrategy to none if it's resulting in a JSON schema not compatible with OpenAI GPT's strict mode.