move app to app/ subdir

app/src/promptConstructor/format.test.ts

@@ -0,0 +1,10 @@
+import { expect, test } from "vitest";
+import { stripTypes } from "./format";
+
+test("stripTypes", () => {
+  expect(stripTypes(`const foo: string = "bar";`)).toBe(`const foo = "bar";`);
+});
+
+test("stripTypes with invalid syntax", () => {
+  expect(stripTypes(`asdf foo: string = "bar"`)).toBe(`asdf foo: string = "bar"`);
+});

app/src/promptConstructor/format.ts

@@ -0,0 +1,31 @@
+import prettier from "prettier/standalone";
+import parserTypescript from "prettier/plugins/typescript";
+
+// @ts-expect-error for some reason missing from types
+import parserEstree from "prettier/plugins/estree";
+
+import * as babel from "@babel/standalone";
+
+export function stripTypes(tsCode: string): string {
+  const options = {
+    presets: ["typescript"],
+    filename: "file.ts",
+  };
+
+  try {
+    const result = babel.transform(tsCode, options);
+    return result.code ?? tsCode;
+  } catch (error) {
+    // console.error("Error stripping types", error);
+    return tsCode;
+  }
+}
+
+export default async function formatPromptConstructor(code: string): Promise<string> {
+  return await prettier.format(stripTypes(code), {
+    parser: "typescript",
+    plugins: [parserTypescript, parserEstree],
+    // We're showing these in pretty narrow panes so let's keep the print width low
+    printWidth: 60,
+  });
+}

app/src/promptConstructor/migrate.test.ts

@@ -0,0 +1,56 @@
+import "dotenv/config";
+import dedent from "dedent";
+import { expect, test } from "vitest";
+import { migrate1to2, migrate2to3 } from "./migrate";
+
+test("migrate1to2", () => {
+  const promptConstructor = dedent`
+  // Test comment
+
+  prompt = {
+    model: "gpt-3.5-turbo-0613",
+    messages: [
+      {
+        role: "user",
+        content: "What is the capital of China?"
+      }
+    ]
+  }
+  `;
+
+  const migrated = migrate1to2(promptConstructor);
+  expect(migrated).toBe(dedent`
+    // Test comment
+    
+    definePrompt("openai/ChatCompletion", {
+      model: "gpt-3.5-turbo-0613",
+      messages: [
+        {
+          role: "user",
+          content: "What is the capital of China?"
+        }
+      ]
+    })
+  `);
+});
+
+test("migrate2to3", () => {
+  const promptConstructor = dedent`
+  // Test comment
+
+  definePrompt("anthropic", {
+    model: "claude-2.0",
+    prompt: "What is the capital of China?"
+  })
+  `;
+
+  const migrated = migrate2to3(promptConstructor);
+  expect(migrated).toBe(dedent`
+    // Test comment
+    
+    definePrompt("anthropic/completion", {
+      model: "claude-2.0",
+      prompt: "What is the capital of China?"
+    })
+  `);
+});

app/src/promptConstructor/migrate.ts

@@ -0,0 +1,125 @@
+import "dotenv/config";
+import * as recast from "recast";
+import { type ASTNode } from "ast-types";
+import { fileURLToPath } from "url";
+import parsePromptConstructor from "./parse";
+import { prisma } from "~/server/db";
+import { promptConstructorVersion } from "./version";
+const { builders: b } = recast.types;
+
+export const migrate1to2 = (fnBody: string): string => {
+  const ast: ASTNode = recast.parse(fnBody);
+
+  recast.visit(ast, {
+    visitAssignmentExpression(path) {
+      const node = path.node;
+      if ("name" in node.left && node.left.name === "prompt") {
+        const functionCall = b.callExpression(b.identifier("definePrompt"), [
+          b.literal("openai/ChatCompletion"),
+          node.right,
+        ]);
+        path.replace(functionCall);
+      }
+      return false;
+    },
+  });
+
+  return recast.print(ast).code;
+};
+
+export const migrate2to3 = (fnBody: string): string => {
+  const ast: ASTNode = recast.parse(fnBody);
+
+  recast.visit(ast, {
+    visitCallExpression(path) {
+      const node = path.node;
+
+      // Check if the function being called is 'definePrompt'
+      if (
+        recast.types.namedTypes.Identifier.check(node.callee) &&
+        node.callee.name === "definePrompt" &&
+        node.arguments.length > 0 &&
+        recast.types.namedTypes.Literal.check(node.arguments[0]) &&
+        node.arguments[0].value === "anthropic"
+      ) {
+        node.arguments[0].value = "anthropic/completion";
+      }
+
+      return false;
+    },
+  });
+
+  return recast.print(ast).code;
+};
+
+const migrations: Record<number, (fnBody: string) => string> = {
+  2: migrate1to2,
+  3: migrate2to3,
+};
+
+const applyMigrations = (
+  promptConstructor: string,
+  currentVersion: number,
+  targetVersion: number,
+) => {
+  let migratedFn = promptConstructor;
+
+  for (let v = currentVersion + 1; v <= targetVersion; v++) {
+    const migrationFn = migrations[v];
+    if (migrationFn) {
+      migratedFn = migrationFn(migratedFn);
+    }
+  }
+
+  return migratedFn;
+};
+
+export default async function migrateConstructFns(targetVersion: number) {
+  const prompts = await prisma.promptVariant.findMany({
+    where: { promptConstructorVersion: { lt: targetVersion } },
+  });
+  console.log(`Migrating ${prompts.length} prompts to version ${targetVersion}`);
+  await Promise.all(
+    prompts.map(async (variant) => {
+      const currentVersion = variant.promptConstructorVersion;
+
+      try {
+        const migratedFn = applyMigrations(
+          variant.promptConstructor,
+          currentVersion,
+          targetVersion,
+        );
+
+        const parsedFn = await parsePromptConstructor(migratedFn);
+        if ("error" in parsedFn) {
+          throw new Error(parsedFn.error);
+        }
+        await prisma.promptVariant.update({
+          where: {
+            id: variant.id,
+          },
+          data: {
+            promptConstructor: migratedFn,
+            promptConstructorVersion: targetVersion,
+            modelProvider: parsedFn.modelProvider,
+            model: parsedFn.model,
+          },
+        });
+      } catch (e) {
+        console.error("Error migrating promptConstructor for variant", variant.id, e);
+      }
+    }),
+  );
+}
+
+// If we're running this file directly, run the migration to the latest version
+if (process.argv.at(-1) === fileURLToPath(import.meta.url)) {
+  const latestVersion = Math.max(...Object.keys(migrations).map(Number));
+  if (latestVersion !== promptConstructorVersion) {
+    throw new Error(
+      `The latest migration is ${latestVersion}, but the promptConstructorVersion is ${promptConstructorVersion}`,
+    );
+  }
+  await migrateConstructFns(promptConstructorVersion);
+  console.log("Done");
+}

app/src/promptConstructor/parse.test.ts

@@ -0,0 +1,45 @@
+import { expect, test } from "vitest";
+import parsePromptConstructor from "./parse";
+import assert from "assert";
+
+// Note: this has to be run with `vitest --no-threads` option or else
+// isolated-vm seems to throw errors
+test("parsePromptConstructor", async () => {
+  const constructed = await parsePromptConstructor(
+    `
+    // These sometimes have a comment
+
+    definePrompt("openai/ChatCompletion", {
+      model: "gpt-3.5-turbo-0613",
+      messages: [
+        {
+          role: "user",
+          content: \`What is the capital of \${scenario.country}?\`
+        }
+      ]
+    })
+    `,
+    { country: "Bolivia" },
+  );
+
+  expect(constructed).toEqual({
+    modelProvider: "openai/ChatCompletion",
+    model: "gpt-3.5-turbo-0613",
+    modelInput: {
+      messages: [
+        {
+          content: "What is the capital of Bolivia?",
+          role: "user",
+        },
+      ],
+      model: "gpt-3.5-turbo-0613",
+    },
+  });
+});
+
+test("bad syntax", async () => {
+  const parsed = await parsePromptConstructor(`definePrompt("openai/ChatCompletion", {`);
+
+  assert("error" in parsed);
+  expect(parsed.error).toContain("Unexpected end of input");
+});

app/src/promptConstructor/parse.ts

@@ -0,0 +1,92 @@
+import modelProviders from "~/modelProviders/modelProviders";
+import ivm from "isolated-vm";
+import { isObject, isString } from "lodash-es";
+import { type JsonObject } from "type-fest";
+import { validate } from "jsonschema";
+
+export type ParsedPromptConstructor<T extends keyof typeof modelProviders> = {
+  modelProvider: T;
+  model: keyof (typeof modelProviders)[T]["models"];
+  modelInput: Parameters<(typeof modelProviders)[T]["getModel"]>[0];
+};
+
+const isolate = new ivm.Isolate({ memoryLimit: 128 });
+
+export default async function parsePromptConstructor(
+  promptConstructor: string,
+  scenario: JsonObject | undefined = {},
+): Promise<ParsedPromptConstructor<keyof typeof modelProviders> | { error: string }> {
+  try {
+    const modifiedConstructFn = promptConstructor.replace(
+      "definePrompt(",
+      "global.prompt = definePrompt(",
+    );
+
+    const code = `
+    const scenario = ${JSON.stringify(scenario ?? {}, null, 2)};
+    
+    const definePrompt = (modelProvider, input) => ({
+      modelProvider,
+      input
+    })
+
+    ${modifiedConstructFn}
+  `;
+
+    const context = await isolate.createContext();
+    const jail = context.global;
+    await jail.set("global", jail.derefInto());
+
+    const script = await isolate.compileScript(code);
+    await script.run(context);
+    const promptReference = (await context.global.get("prompt")) as ivm.Reference;
+    const prompt = await promptReference.copy();
+
+    if (!isObject(prompt)) {
+      return { error: "definePrompt did not return an object" };
+    }
+    if (!("modelProvider" in prompt) || !isString(prompt.modelProvider)) {
+      return { error: "definePrompt did not return a valid modelProvider" };
+    }
+
+    const provider =
+      prompt.modelProvider in modelProviders &&
+      modelProviders[prompt.modelProvider as keyof typeof modelProviders];
+    if (!provider) {
+      return { error: "definePrompt did not return a known modelProvider" };
+    }
+    if (!("input" in prompt) || !isObject(prompt.input)) {
+      return { error: "definePrompt did not return an input" };
+    }
+
+    const validationResult = validate(prompt.input, provider.inputSchema);
+    if (!validationResult.valid)
+      return {
+        error: `definePrompt did not return a valid input: ${validationResult.errors
+          .map((e) => e.stack)
+          .join(", ")}`,
+      };
+
+    // We've validated the JSON schema so this should be safe
+    const input = prompt.input as Parameters<(typeof provider)["getModel"]>[0];
+
+    const model = provider.getModel(input);
+    if (!model) {
+      return {
+        error: `definePrompt did not return a known model for the provider ${prompt.modelProvider}`,
+      };
+    }
+
+    return {
+      modelProvider: prompt.modelProvider as keyof typeof modelProviders,
+      model,
+      modelInput: input,
+    };
+  } catch (e) {
+    const msg =
+      isObject(e) && "message" in e && isString(e.message)
+        ? e.message
+        : "unknown error parsing definePrompt script";
+    return { error: msg };
+  }
+}

app/src/promptConstructor/version.ts

@@ -0,0 +1 @@
+export const promptConstructorVersion = 3;