lib/ask-ai/ask-ai-with-previous-attempts.ts

@@ -1,4 +1,5 @@
 import { openai } from "lib/ai/openai"
+import { getPrimarySourceCodeFromVfs } from "lib/utils/get-primary-source-code-from-vfs"
 
 interface AttemptHistory {
   code: string
@@ -8,12 +9,14 @@ interface AttemptHistory {
 export const askAiWithPreviousAttempts = async ({
   prompt,
   systemPrompt,
+  vfs,
   previousAttempts,
   onStream,
   openaiClient,
 }: {
   prompt: string
   systemPrompt: string
+  vfs?: Record<string, string>
   previousAttempts?: AttemptHistory[]
   onStream?: (chunk: string) => void
   openaiClient?: typeof openai
@@ -25,6 +28,19 @@ export const askAiWithPreviousAttempts = async ({
       { role: "user", content: prompt },
     ]
 
+  const primarySourceCode = getPrimarySourceCodeFromVfs(vfs)
+
+  if (primarySourceCode) {
+    messages.push({
+      role: "assistant",
+      content: "Please modify the code provided by the user.",
+    })
+    messages.push({
+      role: "user",
+      content: primarySourceCode,
+    })
+  }
+
   if (previousAttempts?.length) {
     messages.push({
       role: "user",

lib/tscircuit-coder/run-ai-with-error-correction.ts

@@ -2,6 +2,7 @@ import { askAiWithPreviousAttempts } from "../ask-ai/ask-ai-with-previous-attemp
 import { saveAttemptLog } from "lib/utils/save-attempt"
 import type OpenAI from "openai"
 import { evaluateTscircuitCode } from "../utils/evaluate-tscircuit-code"
+import { getPrimarySourceCodeFromVfs } from "lib/utils/get-primary-source-code-from-vfs"
 
 const createAttemptFile = ({
   fileName,
@@ -68,6 +69,7 @@ export const runAiWithErrorCorrection = async ({
   const aiResponse = await askAiWithPreviousAttempts({
     prompt,
     systemPrompt,
+    vfs,
     previousAttempts,
     onStream,
     openaiClient,

lib/tscircuit-coder/tscircuitCoder.ts

@@ -11,10 +11,13 @@ export interface TscircuitCoderEvents {
 export interface TscircuitCoder {
   vfs: { [filepath: string]: string }
   availableOptions: { name: string; options: string[] }[]
-  submitPrompt: (
-    prompt: string,
-    options?: { selectedMicrocontroller?: string },
-  ) => Promise<void>
+  submitPrompt: ({
+    prompt,
+    options,
+  }: {
+    prompt: string
+    options?: { selectedMicrocontroller?: string }
+  }) => Promise<void>
   on<K extends keyof TscircuitCoderEvents>(
     event: K,
     listener: (payload: TscircuitCoderEvents[K]) => void,
@@ -35,31 +38,37 @@ export class TscircuitCoderImpl extends EventEmitter implements TscircuitCoder {
     this.openaiClient = openaiClient
   }
 
-  async submitPrompt(
-    prompt: string,
-    options?: { selectedMicrocontroller?: string },
-  ): Promise<void> {
+  async submitPrompt({
+    prompt,
+    options,
+  }: {
+    prompt: string
+    options?: { selectedMicrocontroller?: string }
+  }): Promise<void> {
     const systemPrompt = await createLocalCircuitPrompt()
     const promptNumber = Date.now()
     let currentAttempt = ""
     let streamStarted = false
+    const onStream = (chunk: string) => {
+      if (!streamStarted) {
+        this.emit("streamedChunk", "Creating a tscircuit local circuit...")
+        streamStarted = true
+      }
+      currentAttempt += chunk
+      this.emit("streamedChunk", chunk)
+    }
+    const onVfsChanged = () => {
+      this.emit("vfsChanged")
+    }
+
     const result = await runAiWithErrorCorrection({
       prompt,
       systemPrompt,
       promptNumber,
       maxAttempts: 4,
       previousAttempts: [],
-      onStream: (chunk: string) => {
-        if (!streamStarted) {
-          this.emit("streamedChunk", "Creating a tscircuit local circuit...")
-          streamStarted = true
-        }
-        currentAttempt += chunk
-        this.emit("streamedChunk", chunk)
-      },
-      onVfsChanged: () => {
-        this.emit("vfsChanged")
-      },
+      onStream,
+      onVfsChanged,
       vfs: this.vfs,
     })
     if (result.code) {

lib/utils/get-primary-source-code-from-vfs.ts

@@ -0,0 +1,12 @@
+export function getPrimarySourceCodeFromVfs(
+  vfs: Record<string, string> | undefined,
+): string | undefined {
+  let code = undefined
+  if (vfs && Object.keys(vfs).length !== 0) {
+    const finalResultKey = Object.keys(vfs)[Object.keys(vfs).length - 1]
+    if (finalResultKey.includes("final")) {
+      code = vfs[finalResultKey]
+    }
+  }
+  return code
+}

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tscircuit/prompt-benchmarks",
-  "version": "0.0.42",
+  "version": "0.0.43",
   "main": "dist/index.js",
   "repository": "github:tscircuit/prompt-benchmarks",
   "homepage": "https://github.com/tscircuit/prompt-benchmarks#readme",

tests/aiCoder.test.ts → tests/tscircuitCoder.test.ts

@@ -1,5 +1,6 @@
 import { createTscircuitCoder } from "lib/tscircuit-coder/tscircuitCoder"
 import { expect, test } from "bun:test"
+import { getPrimarySourceCodeFromVfs } from "lib/utils/get-primary-source-code-from-vfs"
 
 test("TscircuitCoder submitPrompt streams and updates vfs", async () => {
   const streamedChunks: string[] = []
@@ -12,9 +13,24 @@ test("TscircuitCoder submitPrompt streams and updates vfs", async () => {
     vfsUpdated = true
   })
 
-  await tscircuitCoder.submitPrompt(
-    "create a random complicated circuit that does something cool",
-  )
+  await tscircuitCoder.submitPrompt({
+    prompt: "create bridge rectifier circuit",
+  })
+
+  await tscircuitCoder.submitPrompt({
+    prompt: "add a transistor component",
+  })
+
+  let codeWithTransistor = getPrimarySourceCodeFromVfs(tscircuitCoder.vfs)
+  expect(codeWithTransistor).toInclude("transistor")
+
+  await tscircuitCoder.submitPrompt({
+    prompt: "add a tssop20 chip",
+  })
+
+  let codeWithChip = getPrimarySourceCodeFromVfs(tscircuitCoder.vfs)
+  expect(codeWithChip).toInclude("tssop20")
+  expect(codeWithChip).toInclude("transistor")
 
   expect(streamedChunks.length).toBeGreaterThan(0)
   const vfsKeys = Object.keys(tscircuitCoder.vfs)