structuredCompletion.ts

import OpenAI from "openai";
import { type CompletionCreateParams } from "openai/resources/chat";
import stringify from "json-stringify-pretty-compact";
import type { FromSchema, JSONSchema } from "json-schema-to-ts";

import Ajv from "ajv";

const ajv = new Ajv({ strict: false });

export const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });

// Returns a strongly-typed response from OpenAI's function completion API.
// Uses AJV to validate the response, and `json-schema-to-ts` to correctly
// type it.
// 
// Usage:
// const sentences = structuredCompletion(
//   {
//     model: "gpt-3.5-turbo-0613",
//     messages: [
//       {
//         role: "system",
//         content: "return 10 random sentences",
//       },
//     ],
//   },
//   "output",
//   {
//     type: "object",
//     properties: {
//       sentences: {
//         type: "array",
//         items: {
//           type: "string",
//         },
//       },
//     },
//     required: ["sentences"],
//   } as const
// );
export const structuredCompletion: <T extends JSONSchema>(
  params: CompletionCreateParams.CreateChatCompletionRequestNonStreaming,
  functionName: string,
  functionParams: T
) => Promise<FromSchema<T & { additionalProperties: false }>> = async (
  params,
  functionName,
  functionParams
) => {
  const newParams: CompletionCreateParams.CreateChatCompletionRequestNonStreaming = {
    ...params,
    function_call: {
      name: functionName,
    },
    functions: [
      ...(params.functions ?? []),
      {
        name: functionName,
        parameters: functionParams as Record<string, unknown>,
      },
    ],
  };

  const resp = await openai.chat.completions.create(newParams);

  const msg = resp.choices[0].message;
  if (!msg) throw new Error("No message in completion response");

  // Check to make sure OpenAI actually called the function
  if (!msg.function_call) throw new Error("No function call in completion response");
  let parsedArgs: unknown;
  try {
    // Attempt to parse the arguments as JSON
    parsedArgs = JSON.parse(msg.function_call.arguments ?? "");
  } catch (e) {
    throw new Error(`Failed to parse JSON: ${e}\n${msg.function_call.arguments}`);
  }

  // Ensure the arguments actually match the expected schema
  if (!ajv.validate(functionParams, parsedArgs)) {
    throw new Error(`Invalid response: ${ajv.errorsText()}\n${stringify(parsedArgs)}`);
  }

  // These won't actually be `any` from the caller's perspective since we
  // know they match the schema and define the function return type
  return parsedArgs as any;
};