feat: chain of thought

src/helpers/interactive-mode.ts

@@ -6,7 +6,6 @@ import { getFileSuggestion, getSimpleCompletion } from './llm';
 import { getConfig, setConfigs } from './config';
 import { readFile, writeFile } from 'fs/promises';
 import dedent from 'dedent';
-import { removeBackticks } from './remove-backticks';
 import { formatMessage } from './test';
 import { gray, green } from 'kolorist';
 import { exitOnCancel } from './exit-on-cancel';
@@ -77,21 +76,31 @@ export async function interactiveMode(options: Partial<RunOptions>) {
 
   const testFilePath = filePath.replace(/.(\w+)$/, '.test.$1');
 
-  let testContents = removeBackticks(
+  let testContents = getCodeBlock(
     (await getSimpleCompletion({
       onChunk: (chunk) => {
         process.stderr.write(formatMessage(chunk));
       },
       messages: [
         {
           role: 'system',
-          content:
-            'You return code for a unit test only. No other words, just the code',
+          content: dedent`
+      You are an AI assistant that given a user prompt, returns a markdown for a unit test.
+      1. Think step by step before emiting any code. Think about the shape of the input and output, the behavior and special situations that are relevant to the algorithm.
+
+      2. After planning, return a code block with the test code.
+        - Start with the most basic test case and progress to more complex ones.
+        - Start with the happy path, then edge cases.
+        - Inputs that are invalid, and likely to break the algorithm.
+        - Keep the individual tests small and focused.
+        - Focus in behavior, not implementation.
+
+        Stop emitting after the code block.`,
         },
         {
           role: 'user',
           content: dedent`
-          Please give me a unit test file (can be multiple tests) for the following prompt:
+          Please prepare a unit test file (can be multiple tests) for the following prompt:
           <prompt>
           ${prompt}
           </prompt>
@@ -118,7 +127,7 @@ export async function interactiveMode(options: Partial<RunOptions>) {
         },
       ],
     }))!
-  );
+  )!;
 
   const result = exitOnCancel(
     await text({
@@ -171,3 +180,22 @@ export async function interactiveMode(options: Partial<RunOptions>) {
     lastRunError: '',
   });
 }
+
+export function getCodeBlock(output: string, markdownLang?: string) {
+  const foundCode = output.indexOf('```' + (markdownLang ?? ''));
+  if (foundCode > -1) {
+    const start = output.indexOf('\n', foundCode);
+    if (start === -1) {
+      return undefined;
+    }
+    const end = output.indexOf('```', start);
+    if (end === -1) {
+      console.error('Code block end not found');
+    }
+    return output.substring(start, end === -1 ? undefined : end).trim();
+  }
+  if (markdownLang && foundCode === -1) {
+    return undefined;
+  }
+  return output;
+}

src/helpers/iterate-on-test.ts

@@ -1,10 +1,10 @@
 import { text } from '@clack/prompts';
 import { exitOnCancel } from './exit-on-cancel';
 import { RunOptions } from './run';
-import { removeBackticks } from './remove-backticks';
 import { getSimpleCompletion } from './llm';
 import { formatMessage } from './test';
 import dedent from 'dedent';
+import { getCodeBlock } from './interactive-mode';
 
 export async function iterateOnTest({
   testCode,
@@ -16,7 +16,7 @@ export async function iterateOnTest({
   options: Partial<RunOptions>;
 }) {
   process.stderr.write(formatMessage('\n'));
-  let testContents = removeBackticks(
+  let testContents = getCodeBlock(
     (await getSimpleCompletion({
       onChunk: (chunk) => {
         process.stderr.write(formatMessage(chunk));
@@ -35,7 +35,7 @@ export async function iterateOnTest({
           ${options.prompt}
           </prompt>
 
-          The test will be located at \`${options.testFile}\` and the code to test will be located at 
+          The test will be located at \`${options.testFile}\` and the code to test will be located at
           \`${options.outputFile}\`.
 
           The current test code is:
@@ -56,7 +56,7 @@ export async function iterateOnTest({
         },
       ],
     }))!
-  );
+  )!;
   console.log(formatMessage('\n'));
 
   const result = exitOnCancel(

src/helpers/llm.ts

@@ -147,6 +147,8 @@ export const getSimpleCompletion = async function (options: {
   const completion = await openai.chat.completions.create({
     model: model || defaultModel,
     messages: options.messages,
+    temperature: 0,
+    seed: 42,
     stream: true,
   });
 

src/helpers/systemPrompt.ts

@@ -1,2 +1,10 @@
-export const systemPrompt =
-  'You take a prompt and test and generate code accordingly. You only output code and nothing else - your output just a code string, like "const hello = \'world\'", not markdown (aka do NOT put three backticks around the code). Be sure your code exports function that can be called by an external test file. Make sure your code is reusable and not overly hardcoded to match the promt. Use two spaces for indents. Add logs if helpful for debugging, you will get the log output on your next try to help you debug. Always return a complete code snippet that can execute, nothing partial and never say "rest of your code" or similar, I will copy and paste your code into my file without modification, so it cannot have gaps or parts where you say to put the "rest of the code" back in. Think things through first - write out your thoughts and then write the code.';
+export const systemPrompt = `You take a prompt and unit tests and generate a function accordingly.
+
+1. Think step by step about the algorithm, reasoning about the problem and the solution, similar algirithms, the data structures, state and strategy you will use. Do not write any code in this step.
+
+2. Emit a markdown code block with the generated code (function that satisfies all the tests and the prompt).
+ - Be sure your code exports function that can be called by an external test file.
+ - Make sure your code is reusable and not overly hardcoded to match the promt.
+ - Use two spaces for indents. Add logs if helpful for debugging, you will get the log output on your next try to help you debug.
+ - Always return a complete code snippet that can execute, nothing partial and never say "rest of your code" or similar, I will copy and paste your code into my file without modification, so it cannot have gaps or parts where you say to put the "rest of the code" back in.
+ - Only emit the function, not the tests.`;