Skip to content

Commit

Permalink
fix(helpers/zod): correct schema generation for recursive schemas
Browse files Browse the repository at this point in the history
  • Loading branch information
@RobertCraigie @stainless-app
RobertCraigie authored and stainless-app[bot] committed Aug 7, 2024
1 parent 086421f commit 8b1b4b3
Show file tree
Hide file tree
Showing 3 changed files with 259 additions and 4 deletions.
12 changes: 8 additions & 4 deletions src/helpers/zod.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,12 @@ import {
} from '../lib/parser';
import { zodToJsonSchema as _zodToJsonSchema } from '../_vendor/zod-to-json-schema';

function zodToJsonSchema(schema: z.ZodType): Record<string, unknown> {
return _zodToJsonSchema(schema, { openaiStrictMode: true });
function zodToJsonSchema(schema: z.ZodType, options: { name: string }): Record<string, unknown> {
return _zodToJsonSchema(schema, {
openaiStrictMode: true,
name: options.name,
nameStrategy: 'duplicate-ref',
});
}

/**
Expand Down Expand Up @@ -61,7 +65,7 @@ export function zodResponseFormat<ZodInput extends z.ZodType>(
...props,
name,
strict: true,
schema: zodToJsonSchema(zodObject),
schema: zodToJsonSchema(zodObject, { name }),
},
},
(content) => zodObject.parse(JSON.parse(content)),
Expand Down Expand Up @@ -89,7 +93,7 @@ export function zodFunction<Parameters extends z.ZodType>(options: {
type: 'function',
function: {
name: options.name,
parameters: zodToJsonSchema(options.parameters),
parameters: zodToJsonSchema(options.parameters, { name: options.name }),
strict: true,
...(options.description ? { description: options.description } : undefined),
},
Expand Down
28 changes: 28 additions & 0 deletions tests/lib/__snapshots__/parser.test.ts.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,3 +27,31 @@ exports[`.parse() zod deserialises response_format 1`] = `
}
"
`;

exports[`.parse() zod top-level recursive schemas 1`] = `
"{
"id": "chatcmpl-9taiMDrRVRIkk1Xg1yE82UjnYuZjt",
"object": "chat.completion",
"created": 1723036198,
"model": "gpt-4o-2024-08-06",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "{\\"type\\":\\"form\\",\\"label\\":\\"User Profile Form\\",\\"children\\":[{\\"type\\":\\"field\\",\\"label\\":\\"Full Name\\",\\"children\\":[],\\"attributes\\":[{\\"name\\":\\"type\\",\\"value\\":\\"text\\"},{\\"name\\":\\"placeholder\\",\\"value\\":\\"Enter your full name\\"}]},{\\"type\\":\\"field\\",\\"label\\":\\"Email Address\\",\\"children\\":[],\\"attributes\\":[{\\"name\\":\\"type\\",\\"value\\":\\"email\\"},{\\"name\\":\\"placeholder\\",\\"value\\":\\"Enter your email address\\"}]},{\\"type\\":\\"field\\",\\"label\\":\\"Phone Number\\",\\"children\\":[],\\"attributes\\":[{\\"name\\":\\"type\\",\\"value\\":\\"tel\\"},{\\"name\\":\\"placeholder\\",\\"value\\":\\"Enter your phone number\\"}]},{\\"type\\":\\"button\\",\\"label\\":\\"Submit\\",\\"children\\":[],\\"attributes\\":[{\\"name\\":\\"type\\",\\"value\\":\\"submit\\"}]}],\\"attributes\\":[{\\"name\\":\\"method\\",\\"value\\":\\"post\\"},{\\"name\\":\\"action\\",\\"value\\":\\"/submit-profile\\"}]}",
"refusal": null
},
"logprobs": null,
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 38,
"completion_tokens": 168,
"total_tokens": 206
},
"system_fingerprint": "fp_845eaabc1f"
}
"
`;
223 changes: 223 additions & 0 deletions tests/lib/parser.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,5 +44,228 @@ describe('.parse()', () => {
}
`);
});

test('top-level recursive schemas', async () => {
const UI: any = z.lazy(() =>
z.object({
type: z.enum(['div', 'button', 'header', 'section', 'field', 'form']),
label: z.string(),
children: z.array(UI),
attributes: z.array(
z.object({
name: z.string(),
value: z.string(),
}),
),
}),
);

const completion = await makeSnapshotRequest((openai) =>
openai.beta.chat.completions.parse({
model: 'gpt-4o-2024-08-06',
messages: [
{
role: 'system',
content: 'You are a UI generator AI. Convert the user input into a UI.',
},
{ role: 'user', content: 'Make a User Profile Form with 3 fields' },
],
response_format: zodResponseFormat(UI, 'ui'),
}),
);

expect(completion.choices[0]?.message).toMatchInlineSnapshot(`
{
"content": "{"type":"form","label":"User Profile Form","children":[{"type":"field","label":"Full Name","children":[],"attributes":[{"name":"type","value":"text"},{"name":"placeholder","value":"Enter your full name"}]},{"type":"field","label":"Email Address","children":[],"attributes":[{"name":"type","value":"email"},{"name":"placeholder","value":"Enter your email address"}]},{"type":"field","label":"Phone Number","children":[],"attributes":[{"name":"type","value":"tel"},{"name":"placeholder","value":"Enter your phone number"}]},{"type":"button","label":"Submit","children":[],"attributes":[{"name":"type","value":"submit"}]}],"attributes":[{"name":"method","value":"post"},{"name":"action","value":"/submit-profile"}]}",
"parsed": {
"attributes": [
{
"name": "method",
"value": "post",
},
{
"name": "action",
"value": "/submit-profile",
},
],
"children": [
{
"attributes": [
{
"name": "type",
"value": "text",
},
{
"name": "placeholder",
"value": "Enter your full name",
},
],
"children": [],
"label": "Full Name",
"type": "field",
},
{
"attributes": [
{
"name": "type",
"value": "email",
},
{
"name": "placeholder",
"value": "Enter your email address",
},
],
"children": [],
"label": "Email Address",
"type": "field",
},
{
"attributes": [
{
"name": "type",
"value": "tel",
},
{
"name": "placeholder",
"value": "Enter your phone number",
},
],
"children": [],
"label": "Phone Number",
"type": "field",
},
{
"attributes": [
{
"name": "type",
"value": "submit",
},
],
"children": [],
"label": "Submit",
"type": "button",
},
],
"label": "User Profile Form",
"type": "form",
},
"refusal": null,
"role": "assistant",
"tool_calls": [],
}
`);

expect(zodResponseFormat(UI, 'ui').json_schema).toMatchInlineSnapshot(`
{
"name": "ui",
"schema": {
"$schema": "http://json-schema.org/draft-07/schema#",
"additionalProperties": false,
"definitions": {
"ui": {
"additionalProperties": false,
"properties": {
"attributes": {
"items": {
"additionalProperties": false,
"properties": {
"name": {
"type": "string",
},
"value": {
"type": "string",
},
},
"required": [
"name",
"value",
],
"type": "object",
},
"type": "array",
},
"children": {
"items": {
"$ref": "#/definitions/ui",
},
"type": "array",
},
"label": {
"type": "string",
},
"type": {
"enum": [
"div",
"button",
"header",
"section",
"field",
"form",
],
"type": "string",
},
},
"required": [
"type",
"label",
"children",
"attributes",
],
"type": "object",
},
},
"properties": {
"attributes": {
"items": {
"additionalProperties": false,
"properties": {
"name": {
"type": "string",
},
"value": {
"type": "string",
},
},
"required": [
"name",
"value",
],
"type": "object",
},
"type": "array",
},
"children": {
"items": {
"$ref": "#/definitions/ui",
},
"type": "array",
},
"label": {
"type": "string",
},
"type": {
"enum": [
"div",
"button",
"header",
"section",
"field",
"form",
],
"type": "string",
},
},
"required": [
"type",
"label",
"children",
"attributes",
],
"type": "object",
},
"strict": true,
}
`);
});
});
});

0 comments on commit 8b1b4b3

Please sign in to comment.