refactor runInitialPrompt

tscircuit · Oct 8, 2024 · b75d98e · b75d98e
1 parent 88269a3
commit b75d98e
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 23 deletions.
diff --git a/lib/code-runner/code-runner-context.ts b/lib/code-runner/code-runner-context.ts
@@ -24,3 +24,7 @@ export interface CodeRunnerContext {
    */
   outputType: OutputType
 }
+
+export interface PromptAndRunnerContext
+  extends PromptContext,
+    CodeRunnerContext {}
diff --git a/lib/code-runner/run-prompt.tsx b/lib/code-runner/run-prompt.tsx
@@ -5,22 +5,20 @@ import { safeEvaluateCode } from "lib/code-runner/safe-evaluate-code"
 import { extractCodefence } from "extract-codefence"
 import { anthropic } from "lib/code-runner/anthropic"
 import Debug from "debug"
+import type {
+  PromptAndRunnerContext,
+  PromptContext,
+} from "./code-runner-context"
 
 const debug = Debug("tscircuit:prompt")
 
 export const runInitialPrompt = async (
-  systemPrompt: string,
-  userPrompt: string,
-  opts: {
-    model?: "claude-3-5-sonnet-20240620" | "claude-3-haiku-20240307"
-    type?: "board" | "footprint" | "package" | "model"
-    availableImports?: Record<string, string>
-    preSuppliedImports?: Record<string, any>
-  } = {},
+  { systemPrompt, userPrompt }: { systemPrompt: string; userPrompt: string },
+  context: PromptAndRunnerContext,
 ) => {
-  const type = opts.type ?? "board"
+  const type = context.outputType
   const completion = await anthropic.messages.create({
-    model: opts.model ?? "claude-3-5-sonnet-20240620",
+    model: context.model,
     system: systemPrompt,
     messages: [
       {
@@ -51,7 +49,11 @@ export const runInitialPrompt = async (
     syntaxError,
     circuitErrors,
     typescriptErrors,
-  } = safeEvaluateCode(codefence, type, opts.preSuppliedImports)
+  } = safeEvaluateCode(
+    codefence,
+    context.outputType,
+    context.preSuppliedImports,
+  )
 
   if (success) {
     return {

diff --git a/tests/prompts/create-circuit-board1/create-circuit-board1-sample1.test.ts b/tests/prompts/create-circuit-board1/create-circuit-board1-sample1.test.ts
@@ -6,10 +6,13 @@ import { sample1 } from "tests/samples/sample1"
 test("create-circuit-board1-prompt1", async () => {
   const systemPrompt = createCircuitBoard1Template({ currentCode: "" })
 
-  const { success, circuit } = await runInitialPrompt(systemPrompt, sample1, {
-    model: "claude-3-haiku-20240307",
-    type: "board",
-  })
+  const { success, circuit } = await runInitialPrompt(
+    { systemPrompt, userPrompt: sample1 },
+    {
+      model: "claude-3-haiku-20240307",
+      outputType: "board",
+    }
+  )
 
   expect(success).toBe(true)
 

diff --git a/tests/smoke/create-circuit-board1-sample1.test.ts b/tests/smoke/create-circuit-board1-sample1.test.ts
@@ -8,12 +8,11 @@ test("create-circuit-board1-prompt1", async () => {
   const systemPrompt = createCircuitBoard1Template({ currentCode: "" })
 
   const { success, circuit, codefence } = await runInitialPrompt(
-    systemPrompt,
-    sample1,
+    { systemPrompt, userPrompt: sample1 },
     {
       model: "claude-3-haiku-20240307",
-      type: "board",
-    },
+      outputType: "board",
+    }
   )
 
   expect(success).toBe(true)

diff --git a/tests/smoke/create-circuit-board1-with-import.test.tsx b/tests/smoke/create-circuit-board1-with-import.test.tsx
@@ -21,11 +21,10 @@ test("create-circuit-board1-prompt1", async () => {
     />
   )
   const { success, circuit, codefence, error } = await runInitialPrompt(
-    systemPrompt,
-    "an led with an 0402 footprint",
+    { systemPrompt, userPrompt: "an led with an 0402 footprint" },
     {
       model: "claude-3-haiku-20240307",
-      type: "board",
+      outputType: "board",
       availableImports: {
         "@tsci/seveibar.micro-usb": `
         
@@ -50,7 +49,7 @@ import MicroUsb from "@tsci/seveibar.micro-usb"
       preSuppliedImports: {
         "@tsci/seveibar.micro-usb": MicroUsb,
       },
-    },
+    }
   )
 
   expect(success).toBe(true)