createJSONCompletion.ts

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;
    }
  };
}