🤖

src/helpers/interactive-mode.ts

@@ -130,7 +130,7 @@ export async function interactiveMode(options: Partial<RunOptions>) {
         },
       ],
     }))!
-  )!;
+  );
 
   const result = exitOnCancel(
     await text({
@@ -184,21 +184,18 @@ export async function interactiveMode(options: Partial<RunOptions>) {
   });
 }
 
-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();
+export function getCodeBlock(output: string) {
+  const foundCode = output.indexOf('```');
+  if (foundCode === -1) {
+    return output;
+  }
+  const start = output.indexOf('\n', foundCode);
+  if (start === -1) {
+    return output.slice(foundCode);
   }
-  if (markdownLang && foundCode === -1) {
-    return undefined;
+  const end = output.indexOf('```', start);
+  if (end === -1) {
+    console.error('Code block end not found');
   }
-  return output;
+  return output.slice(start, end === -1 ? undefined : end).trim();
 }

src/helpers/iterate-on-test.ts

@@ -56,7 +56,7 @@ export async function iterateOnTest({
         },
       ],
     }))!
-  )!;
+  );
   console.log(formatMessage('\n'));
 
   const result = exitOnCancel(

src/helpers/llm.ts

@@ -271,7 +271,7 @@ export const getCompletion = async function (options: {
         })
         .on('textDone', () => {
           process.stdout.write('\n');
-          const output = getCodeBlock(result)!;
+          const output = getCodeBlock(result);
           captureLlmRecord(options.messages, output, mockLlmRecordFile);
           resolve(output);
         });

src/helpers/mock-llm.ts

@@ -51,8 +51,13 @@ export const mockedLlmCompletion = async (
   const completion = jsonLlmRecording.completions.find(
     (completion: { inputs: any }) => {
       // Match on system input only
+      const content = completion.inputs[0].content;
+      if (typeof messages[0].content === 'string') {
+        return messages[0].content.includes(content);
+      }
       return (
-        JSON.stringify(completion.inputs[0]) === JSON.stringify(messages[0])
+        JSON.stringify(completion.inputs[0]) ===
+        JSON.stringify(messages[0].content)
       );
     }
   );

test/fixtures/add.json

@@ -4,7 +4,7 @@
       "inputs": [
         {
           "role": "system",
-          "content": "You return a file path only. No other words, just one file path"
+          "content": "You are an assistant that given a snapshot of the current filesystem suggests a relative file path for the code algorithm mentioned in the prompt. No other words, just one file path"
         },
         {
           "role": "user",
@@ -17,11 +17,11 @@
       "inputs": [
         {
           "role": "system",
-          "content": "You return code for a unit test only. No other words, just the code"
+          "content": "You are an AI assistant that given a user prompt"
         },
         {
           "role": "user",
-          "content": "Please give me a unit test file (can be multiple tests) for the following prompt:\n        <prompt>\n        Adds numbers together\n        </prompt>\n\n        The test will be located at `test/add.test.ts` and the code to test will be located at \n        `test/add.ts`.\n\n        Here is a copy of a couple example tests in the repo:\n        <tests>\n        import { test, expect } from 'vitest';\nimport { removeBackticks } from './remove-backticks';\n\n// Remove backticks from a string, for instance to remove the\n// markdown backticks (+ language name) around code returned that\n// should just be that code\ntest('should remove backticks', () => {\nexpect(removeBackticks('```\nhello\n```')).toBe('hello');\nexpect(removeBackticks('```typescript\nhello\nworld\n```')).toBe(\n'hello\nworld'\n);\nexpect(removeBackticks('```js\nhello\nworld\n```')).toBe('hello\nworld');\n});\n\nimport { getOpenAi, getSimpleCompletion } from './llm';\nimport { KnownError } from './error';\nimport { expect, describe, it, vi } from 'vitest';\nimport OpenAI from 'openai';\nimport { ChatCompletionMessageParam } from 'openai/resources';\n\nconst mocks = vi.hoisted(() => {\nreturn {\nopenAIConstructor: vi.fn(),\ngetConfig: vi.fn(),\ncreate: vi.fn(),\n};\n});\n\nvi.mock('./config', () => {\nreturn {\ngetConfig: mocks.getConfig,\n};\n});\n\nvi.mock('openai', () => {\nreturn {\ndefault: mocks.openAIConstructor,\n};\n});\n\nmocks.openAIConstructor.mockImplementation(() => {\nreturn {\nchat: {\n  completions: {\n    create: mocks.create,\n  },\n},\n};\n});\n\nconst defaultConfig = {\nOPENAI_KEY: 'my-openai-key',\nOPENAI_API_ENDPOINT: 'https://api.openai.com/v1',\n};\n\ndescribe('getOpenAi', () => {\nit('should throw a KnownError if OPENAI_KEY is blank', async () => {\nmocks.getConfig\n  .mockResolvedValueOnce({ OPENAI_KEY: '' })\n  .mockResolvedValueOnce({ OPENAI_KEY: '' });\n\nawait expect(getOpenAi()).rejects.toThrow(KnownError);\nawait expect(getOpenAi()).rejects.toThrow(\n  'Missing OpenAI key. Use `micro-agent config` to set it.'\n);\n});\n\nit('should create a new OpenAI instance with the provided key and endpoint', async () => {\nmocks.getConfig.mockResolvedValueOnce(defaultConfig);\n\nawait getOpenAi();\n\nexpect(OpenAI).toHaveBeenCalledWith({\n  apiKey: 'my-openai-key',\n  baseURL: 'https://api.openai.com/v1',\n});\n});\n});\n\ndescribe('getSimpleCompletion', () => {\nit('should call openai.chat.completions.create with the correct parameters', async () => {\nmocks.getConfig\n  .mockResolvedValueOnce(defaultConfig)\n  .mockResolvedValueOnce(defaultConfig);\nmocks.create.mockResolvedValueOnce([]);\n\nconst messages: ChatCompletionMessageParam[] = [\n  { role: 'system', content: 'Hello' },\n];\nawait getSimpleCompletion({ messages });\n\nexpect(mocks.create).toHaveBeenCalledWith({\n  model: 'gpt-4o',\n  messages,\n  stream: true,\n});\n});\n\nit('should concatenate the output from completion chunks', async () => {\nmocks.getConfig\n  .mockResolvedValueOnce(defaultConfig)\n  .mockResolvedValueOnce(defaultConfig);\nmocks.create.mockResolvedValueOnce([\n  { choices: [{ delta: { content: 'Hello' } }] },\n  { choices: [{ delta: { content: ' World' } }] },\n]);\n\nconst messages: ChatCompletionMessageParam[] = [\n  { role: 'system', content: 'Hello' },\n];\nconst output = await getSimpleCompletion({ messages });\n\nexpect(output).toBe('Hello World');\n});\n\nit('should call options.onChunk for each chunk', async () => {\nmocks.getConfig\n  .mockResolvedValueOnce(defaultConfig)\n  .mockResolvedValueOnce(defaultConfig);\nmocks.create.mockResolvedValueOnce([\n  { choices: [{ delta: { content: 'Hello' } }] },\n  { choices: [{ delta: { content: ' World' } }] },\n]);\n\nconst messages: ChatCompletionMessageParam[] = [\n  { role: 'system', content: 'Hello' },\n];\nconst onChunk = vi.fn();\nconst output = await getSimpleCompletion({ messages, onChunk });\n\nexpect(onChunk).toHaveBeenCalledTimes(2);\nexpect(onChunk).toHaveBeenCalledWith('Hello');\nexpect(onChunk).toHaveBeenCalledWith(' World');\n});\n});\n\n        </tests>"
+          "content": "Please prepare a unit test file (can be multiple tests) for the following prompt:\n        <prompt>\n        Adds numbers together\n        </prompt>\n\n        The test will be located at `test/add.test.ts` and the code to test will be located at\n        `test/add.ts`.\n\n        Here is a copy of a couple example tests in the repo:\n        <tests>\n        import { test, expect } from 'vitest';\nimport { removeBackticks } from './remove-backticks';\n\n// Remove backticks from a string, for instance to remove the\n// markdown backticks (+ language name) around code returned that\n// should just be that code\ntest('should remove backticks', () => {\nexpect(removeBackticks('```\nhello\n```')).toBe('hello');\nexpect(removeBackticks('```typescript\nhello\nworld\n```')).toBe(\n'hello\nworld'\n);\nexpect(removeBackticks('```js\nhello\nworld\n```')).toBe('hello\nworld');\n});\n\nimport { getOpenAi, getSimpleCompletion } from './llm';\nimport { KnownError } from './error';\nimport { expect, describe, it, vi } from 'vitest';\nimport OpenAI from 'openai';\nimport { ChatCompletionMessageParam } from 'openai/resources';\n\nconst mocks = vi.hoisted(() => {\nreturn {\nopenAIConstructor: vi.fn(),\ngetConfig: vi.fn(),\ncreate: vi.fn(),\n};\n});\n\nvi.mock('./config', () => {\nreturn {\ngetConfig: mocks.getConfig,\n};\n});\n\nvi.mock('openai', () => {\nreturn {\ndefault: mocks.openAIConstructor,\n};\n});\n\nmocks.openAIConstructor.mockImplementation(() => {\nreturn {\nchat: {\n  completions: {\n    create: mocks.create,\n  },\n},\n};\n});\n\nconst defaultConfig = {\nOPENAI_KEY: 'my-openai-key',\nOPENAI_API_ENDPOINT: 'https://api.openai.com/v1',\n};\n\ndescribe('getOpenAi', () => {\nit('should throw a KnownError if OPENAI_KEY is blank', async () => {\nmocks.getConfig\n  .mockResolvedValueOnce({ OPENAI_KEY: '' })\n  .mockResolvedValueOnce({ OPENAI_KEY: '' });\n\nawait expect(getOpenAi()).rejects.toThrow(KnownError);\nawait expect(getOpenAi()).rejects.toThrow(\n  'Missing OpenAI key. Use `micro-agent config` to set it.'\n);\n});\n\nit('should create a new OpenAI instance with the provided key and endpoint', async () => {\nmocks.getConfig.mockResolvedValueOnce(defaultConfig);\n\nawait getOpenAi();\n\nexpect(OpenAI).toHaveBeenCalledWith({\n  apiKey: 'my-openai-key',\n  baseURL: 'https://api.openai.com/v1',\n});\n});\n});\n\ndescribe('getSimpleCompletion', () => {\nit('should call openai.chat.completions.create with the correct parameters', async () => {\nmocks.getConfig\n  .mockResolvedValueOnce(defaultConfig)\n  .mockResolvedValueOnce(defaultConfig);\nmocks.create.mockResolvedValueOnce([]);\n\nconst messages: ChatCompletionMessageParam[] = [\n  { role: 'system', content: 'Hello' },\n];\nawait getSimpleCompletion({ messages });\n\nexpect(mocks.create).toHaveBeenCalledWith({\n  model: 'gpt-4o',\n  messages,\n  stream: true,\n});\n});\n\nit('should concatenate the output from completion chunks', async () => {\nmocks.getConfig\n  .mockResolvedValueOnce(defaultConfig)\n  .mockResolvedValueOnce(defaultConfig);\nmocks.create.mockResolvedValueOnce([\n  { choices: [{ delta: { content: 'Hello' } }] },\n  { choices: [{ delta: { content: ' World' } }] },\n]);\n\nconst messages: ChatCompletionMessageParam[] = [\n  { role: 'system', content: 'Hello' },\n];\nconst output = await getSimpleCompletion({ messages });\n\nexpect(output).toBe('Hello World');\n});\n\nit('should call options.onChunk for each chunk', async () => {\nmocks.getConfig\n  .mockResolvedValueOnce(defaultConfig)\n  .mockResolvedValueOnce(defaultConfig);\nmocks.create.mockResolvedValueOnce([\n  { choices: [{ delta: { content: 'Hello' } }] },\n  { choices: [{ delta: { content: ' World' } }] },\n]);\n\nconst messages: ChatCompletionMessageParam[] = [\n  { role: 'system', content: 'Hello' },\n];\nconst onChunk = vi.fn();\nconst output = await getSimpleCompletion({ messages, onChunk });\n\nexpect(onChunk).toHaveBeenCalledTimes(2);\nexpect(onChunk).toHaveBeenCalledWith('Hello');\nexpect(onChunk).toHaveBeenCalledWith(' World');\n});\n});\n\n        </tests>"
         }
       ],
       "output": "```typescript\nimport { test, expect } from 'vitest';\nimport { add } from './add';\n\ntest('should add two positive numbers', () => {\n  expect(add(2, 3)).toBe(5);\n});\n\ntest('should add two negative numbers', () => {\n  expect(add(-2, -3)).toBe(-5);\n});\n\ntest('should add a positive and a negative number', () => {\n  expect(add(2, -3)).toBe(-1);\n});\n\ntest('should add zero to a number', () => {\n  expect(add(0, 5)).toBe(5);\n  expect(add(5, 0)).toBe(5);\n});\n\ntest('should add decimal numbers', () => {\n  expect(add(2.5, 3.2)).toBeCloseTo(5.7, 5);\n  expect(add(-2.5, -3.2)).toBeCloseTo(-5.7, 5);\n});\n\ntest('should handle large numbers', () => {\n  expect(add(1e12, 1e12)).toBe(2e12);\n});\n```"
@@ -30,7 +30,7 @@
       "inputs": [
         {
           "role": "system",
-          "content": "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."
+          "content": "You take a prompt and existing unit tests and generate the function"
         },
         {
           "role": "user",