Skip to content

Instantly share code, notes, and snippets.

@aweary
Last active June 10, 2024 20:27
Show Gist options
  • Save aweary/e37e6f374751bd83228120a7546429ae to your computer and use it in GitHub Desktop.
Save aweary/e37e6f374751bd83228120a7546429ae to your computer and use it in GitHub Desktop.
import { z } from "zod";
import { zodToTs, printNode } from "zod-to-ts";
// Replace with your `openai` thing
import { openai } from "../openai.server";
import endent from "endent";
function createJSONCompletion<T extends z.ZodType>({
prompt,
schema_name,
schema,
model,
default: default_value,
example,
}: {
prompt: string | ((content: string) => Promise<string>);
schema: T;
schema_name?: string;
model: "gpt-4" | "gpt-3.5-turbo";
example: z.infer<T>;
default: z.infer<T>;
}): (content: string) => Promise<z.infer<T>> {
const { node } = zodToTs(schema, schema_name);
const ts_type = printNode(node, {});
return async (content: string) => {
let resolved_prompt = "";
if (typeof prompt === "string") {
resolved_prompt = prompt;
} else {
resolved_prompt = await prompt(content);
}
try {
// gpt-3.5-turbo listens to 'user' better than 'system'
const system_role = model === "gpt-4" ? "system" : "user";
const messages: ChatCompletionRequestMessage[] = [
{
role: system_role,
content: endent`You MUST respond only with valid schema compliant JSON and NO other text.`,
},
{
role: system_role,
content: endent`
* ${/* Put your global context here. Like 'You are a Journal AI...' or whatever you're building */}
* ${resolved_prompt}.
* You MUST return the structured data as a JSON object that is compliant with the following TypeScript type:
\`\`\`typescript
${ts_type}
\`\`\`
Return an example response to confirm you understand the schema and requirements.
`,
},
{
role: "assistant",
content: endent`${JSON.stringify(example)}`,
},
{
role: "user",
content,
},
];
const completion = await openai.createChatCompletion({
model,
messages,
max_tokens: 300,
n: 3,
});
for (const { message } of completion.data.choices) {
try {
const parsed = JSON.parse(message?.content ?? "");
return schema.parse(parsed);
} catch (err) {
continue;
}
}
return default_value;
} catch (err) {
return default_value;
}
};
}
@MCYouks
Copy link

MCYouks commented Apr 16, 2023

Thanks for sharing! Can you explain how the example parameter is supposed to be injected in the prompt? Thanks

@aweary
Copy link
Author

aweary commented Apr 16, 2023

@MCYouks it was meant to go in the assistant message instead of default_value, thanks for pointing that out. Fixed

@rob-gordon
Copy link

How is zodToJsonSchema being used?

@aweary
Copy link
Author

aweary commented Apr 17, 2023

@rob-gordon before zod-to-ts I used that to include a JSON schema in the prompt. Just forgot to remove the import here, thanks for the pointing it out! Fixed.

@rob-gordon
Copy link

Got it! I made a version of this a couple weeks ago but I was passing the zod schema and the zod schema as a string 😅
The zodToTs->printNode is brilliant. Thanks for sharing this!!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment