AI cell generation

packages/api/ai/srcbook-generator.mts

@@ -1,24 +1,44 @@
 import { generateText, GenerateTextResult } from 'ai';
+import { type CodeLanguageType, type CellType } from '@srcbook/shared';
+import { type SessionType } from '../types.mjs';
 import { readFileSync } from 'node:fs';
 import Path from 'node:path';
 import { createOpenAI } from '@ai-sdk/openai';
 import { PROMPTS_DIR } from '../constants.mjs';
+import { encode, decodePartial } from '../srcmd.mjs';
 import { getConfig } from '../config.mjs';
 
-const SYSTEM_PROMPT = readFileSync(Path.join(PROMPTS_DIR, 'srcbook-generator.txt'), 'utf-8');
-const MODEL = 'gpt-4o';
+const OPENAI_MODEL = 'gpt-4o';
 
-type NoToolsGenerateTextResult = GenerateTextResult<{}>;
-/*
- * Given a user request, which is free form text describing their intent,
- * generate a srcbook using an LLM.
- *
- * Currently, this uses openAI and the GPT-4o model, and throws if the
- * openAI API key is not set in the settings.
- * In the future, we can parameterize this with different models, to allow
- * users to use different providers like Anthropic or local ones.
+const makeGenerateSrcbookSystemPrompt = () => {
+  return readFileSync(Path.join(PROMPTS_DIR, 'srcbook-generator.txt'), 'utf-8');
+};
+
+const makeGenerateCellSystemPrompt = (language: CodeLanguageType) => {
+  return readFileSync(Path.join(PROMPTS_DIR, `cell-generator-${language}.txt`), 'utf-8');
+};
+
+const makeGenerateCellUserPrompt = (session: SessionType, insertIdx: number, query: string) => {
+  const inlineSrcbookWithPlaceholder = encode(session.cells, session.metadata, {
+    inline: true,
+    insertCellIdx: insertIdx,
+  });
+
+  const prompt = `==== BEGIN SRCBOOK ====
+${inlineSrcbookWithPlaceholder}
+==== END SRCBOOK ====
+
+==== BEGIN USER REQUEST ====
+${query}
+==== END USER REQUEST ====`;
+  return prompt;
+};
+
+/**
+ * Get the OpenAI client and model configuration.
+ * Throws an error if the OpenAI API key is not set in the settings.
  */
-export async function generateSrcbook(query: string): Promise<NoToolsGenerateTextResult> {
+async function getOpenAIModel() {
   const config = await getConfig();
   if (!config.openaiKey) {
     throw new Error('OpenAI API key is not set');
@@ -29,9 +49,24 @@ export async function generateSrcbook(query: string): Promise<NoToolsGenerateTex
     apiKey: config.openaiKey,
   });
 
+  return openai(OPENAI_MODEL);
+}
+
+type NoToolsGenerateTextResult = GenerateTextResult<{}>;
+/*
+ * Given a user request, which is free form text describing their intent,
+ * generate a srcbook using an LLM.
+ *
+ * Currently, this uses openAI and the GPT-4o model, and throws if the
+ * openAI API key is not set in the settings.
+ * In the future, we can parameterize this with different models, to allow
+ * users to use different providers like Anthropic or local ones.
+ */
+export async function generateSrcbook(query: string): Promise<NoToolsGenerateTextResult> {
+  const model = await getOpenAIModel();
   const result = await generateText({
-    model: openai(MODEL),
-    system: SYSTEM_PROMPT,
+    model: model,
+    system: makeGenerateSrcbookSystemPrompt(),
     prompt: query,
   });
 
@@ -41,3 +76,47 @@ export async function generateSrcbook(query: string): Promise<NoToolsGenerateTex
   }
   return result;
 }
+
+type GenerateCellResult = {
+  error: boolean;
+  errors?: string[];
+  cell?: CellType;
+};
+export async function generateCell(
+  query: string,
+  session: SessionType,
+  insertIdx: number,
+): Promise<GenerateCellResult> {
+  const model = await getOpenAIModel();
+
+  const systemPrompt = makeGenerateCellSystemPrompt(session.metadata.language);
+  const userPrompt = makeGenerateCellUserPrompt(session, insertIdx, query);
+  const result = await generateText({
+    model: model,
+    system: systemPrompt,
+    prompt: userPrompt,
+  });
+
+  // TODO, handle 'length' finish reason with sequencing logic.
+  if (result.finishReason !== 'stop') {
+    console.warn('Generated a cell, but finish_reason was not "stop":', result.finishReason);
+  }
+
+  // Parse the result into cells
+  const text = result.text;
+
+  // TODO figure out logging here. It's incredibly valuable to see the data going to and from the LLM
+  // for debugging, but there are considerations around privacy and log size to think about.
+  const decodeResult = decodePartial(text);
+
+  if (decodeResult.error) {
+    return { error: true, errors: decodeResult.errors };
+  } else {
+    const cells = decodeResult.cells;
+    if (cells.length !== 1) {
+      return { error: true, errors: ['Multiple cells generated. Expected only one.'] };
+    } else {
+      return { error: false, cell: decodeResult.cells[0] };
+    }
+  }
+}

packages/api/prompts/cell-generator-javascript.txt

@@ -0,0 +1,128 @@
+===== BEGIN INSTRUCTIONS CONTEXT =====
+
+You are tasked with generating a Srcbook code cell for the user.
+
+A Srcbook is a JavaScript notebook following a markdown-compatible format called .srcmd. It's an interactive and rich way of programming that follows the literate programming idea. 
+
+## Srcbook spec
+
+Structure of a Srcbook:
+0. The language comment: `<!-- srcbook:{"language":"javascript"} -->`
+1. Title cell (heading 1)
+2. Package.json cell, listing deps
+3. N more cells, which are either:
+  a. Markdown cells (GitHub flavored Markdown)
+  b. JavaScript code cells, which have a filename and source content.
+
+The user is already working on an existing Srcbook, and is asking you to create exactly one cell at the given position described below. The cell can be code and markdown. If unspecified lean towards a code cell. 
+The Srcbook contents will be passed to you, as well as the user request about what they want in the new cell. Your job is to write the code cell or the markdown cell for the user, making sure you honor their request while leveraging the context of the rest of the Srcbook.
+Each code cell needs to have a unique filename, as it maps to a file on disk.
+
+Code cells are valid javascript code. They have a unique filename. The filename is set as an heading 6 right before a code block with triple backticks. These backticks denote a code block and specify the language, which is always javascript. Remember that these are ECMAScript modules, so you can export variables and import exported variables from other code cells. For example.
+
+Markdown cells are regular markdown. Just avoid using heading1 and heading6, as those are reserved by the Srcbook spec.
+===== END INSTRUCTIONS CONTEXT ======
+
+===== BEGIN EXAMPLE SRCBOOK =====
+<!-- srcbook:{"language":"javascript"} -->
+
+# Getting started
+
+###### package.json
+
+```json
+{
+  "type": "module",
+  "dependencies": {
+    "random-words": "^2.0.1"
+  }
+}
+```
+
+## What are Srcbooks?
+
+Srcbooks are an interactive way of programming in JavaScript or TypeScript. They are similar to other notebooks like python's [jupyter notebooks](https://jupyter.org/), but unique in their own ways.
+They are based on the [node](https://nodejs.org/en) runtime.
+
+A Srcbook is composed of **cells**. Currently, there are 4 types of cells:
+ 1. **Title cell**: this is "Getting started" above. There is one per Srcbook.
+ 2. **package.json cell**: this is a special cell that manages dependencies for the Srcbook.
+ 3. **markdown cell**: what you're reading is a markdown cell. It allows you to easily express ideas with rich markup, rather than code comments, an idea called [literate programming](https://en.wikipedia.org/wiki/Literate_programming).
+ 4. **code cell**: think of these as JS or TS files. You can run them or export objects to be used in other cells.
+
+###### simple-code.js
+
+```javascript
+// This is a trivial code cell. You can run me by
+// clicking 'Run' or using the shortcut `cmd` + `enter`.
+console.log("Hello, Srcbook!")
+```
+
+## Dependencies
+
+You can add any external node.js-compatible dependency from [npm](https://www.npmjs.com/). Let's look at an example below by importing the `random-words` library.
+
+You'll need to make sure you install dependencies, which you can do by running the `package.json` cell above.
+
+###### generate-random-word.js
+
+```javascript
+import {generate} from 'random-words';
+
+console.log(generate())
+```
+
+## Importing other cells
+
+Behind the scenes, cells are files of JavaScript or TypeScript code. They are ECMAScript 6 modules. Therefore you can export variables from one file and import them in another.
+
+###### star-wars.js
+
+```javascript
+export const func = (name) => `I am your father, ${name}`
+```
+
+###### logger.js
+
+```javascript
+import {func} from './star-wars.js';
+
+console.log(func("Luke"));
+```
+
+## Using secrets
+
+For security purposes, you should avoid pasting secrets directly into Srcbooks. The mechanism you should leverage is [secrets](/secrets). These are stored securely and are accessed at runtime as environment variables.
+
+Secrets can then be imported in Srcbooks using `process.env.SECRET_NAME`:
+```
+const API_KEY = process.env.SECRET_API_KEY;
+const token = auth(API_KEY);
+```
+===== END EXAMPLE SRCBOOK =====
+
+===== BEGIN FINAL INSTRUCTIONS =====
+The user's Srcbook will be passed to you, surrounded with "==== BEGIN SRCBOOK ====" and "==== END SRCBOOK ====".
+The location of the cell you're providing code for will be marked with  "==== INTRODUCE CELL HERE ====".
+The user's request will be passed to you between "==== BEGIN USER REQUEST ====" and "==== END USER REQUEST ====".
+Your job is to write exactly one cell: the filename and the javascript code for this cell according to the Srcbook spec, or a Markdown cell. Lean towards a code cell if the user request is unclear.
+ONLY RETURN THESE THINGS, NO PREAMBULE, NO SUFFIX, ONLY THE CELL CONTENTS.
+
+Below is an example return value for a code cell that you would return. You would return _only_ what is within the <example> tags:
+<example>
+###### simple-ws-client.js
+```javascript
+import WebSocket from 'ws';
+
+// Reference the same port the server is running on
+const ws = new WebSocket('ws://localhost:5405');
+
+ws.on('open', () => {
+  ws.send('Hello from simple-client.js');
+  ws.close();
+});
+```
+</example>
+
+ Write the best possible code you can, as if you were an expert JavaScript engineer. Focus on being elegant, concise, clear. Keep things simple and explicit, but the user's request is your top priority.
+===== END FINAL INSTRUCTIONS ===