Ai iterate code cell (#144)

* Add basic styles for generation of cells. Add WS req/resp

* MVP working

* Progress before trying codemirror-merge

* MVP working diff experience

* Cleanup frontend code

* Address some PR comments, naming and comment updates

* Add prompt tips

* Put the comments at the top

packages/api/ai/srcbook-generator.mts → packages/api/ai/generate.mts

@@ -2,6 +2,7 @@ import { generateText, GenerateTextResult } from 'ai';
 import {
   type CodeLanguageType,
   type CellType,
+  type CodeCellType,
   randomid,
   type CellWithPlaceholderType,
 } from '@srcbook/shared';
@@ -47,6 +48,32 @@ ${query}
   return prompt;
 };
 
+const makeGenerateCellEditSystemPrompt = (language: CodeLanguageType) => {
+  return readFileSync(Path.join(PROMPTS_DIR, `code-updater-${language}.txt`), 'utf-8');
+};
+
+const makeGenerateCellEditUserPrompt = (
+  query: string,
+  session: SessionType,
+  cell: CodeCellType,
+) => {
+  const inlineSrcbook = encode(session.cells, session.metadata, { inline: true });
+
+  const prompt = `==== BEGIN SRCBOOK ====
+${inlineSrcbook}
+==== END SRCBOOK ====
+
+==== BEGIN CODE CELL ====
+${cell.source}
+==== END CODE CELL ====
+
+==== 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.
@@ -90,16 +117,16 @@ export async function generateSrcbook(query: string): Promise<NoToolsGenerateTex
   return result;
 }
 
-type GenerateCellResult = {
+type GenerateCellsResult = {
   error: boolean;
   errors?: string[];
   cells?: CellType[];
 };
-export async function generateCell(
+export async function generateCells(
   query: string,
   session: SessionType,
   insertIdx: number,
-): Promise<GenerateCellResult> {
+): Promise<GenerateCellsResult> {
   const model = await getOpenAIModel();
 
   const systemPrompt = makeGenerateCellSystemPrompt(session.metadata.language);
@@ -116,15 +143,27 @@ export async function generateCell(
   }
 
   // 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 = decodeCells(text);
+  // TODO: figure out logging.
+  // Data is incredibly valuable for product improvements, but privacy needs to be considered.
+  const decodeResult = decodeCells(result.text);
 
   if (decodeResult.error) {
     return { error: true, errors: decodeResult.errors };
   } else {
     return { error: false, cells: decodeResult.cells };
   }
 }
+
+export async function generateCellEdit(query: string, session: SessionType, cell: CodeCellType) {
+  const model = await getOpenAIModel();
+
+  const systemPrompt = makeGenerateCellEditSystemPrompt(session.metadata.language);
+  const userPrompt = makeGenerateCellEditUserPrompt(query, session, cell);
+  const result = await generateText({
+    model: model,
+    system: systemPrompt,
+    prompt: userPrompt,
+  });
+
+  return result.text;
+}

packages/api/prompts/code-updater-javascript.txt

@@ -0,0 +1,106 @@
+===== BEGIN INSTRUCTIONS CONTEXT =====
+
+You are tasked with editing a code snippet (or "cell") in a Srcbook."
+
+A Srcbook is a JavaScript notebook following a markdown-compatible format.
+
+## 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 edit a specific code cell. 
+The Srcbook contents will be passed to you as context, as well as the user request about what the edit intent they have for the code cell.
+===== 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. 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 specific code cell they want updated will also be passed to you, surrounded with "==== BEGIN CODE CELL ====" and "==== END CODE CELL ====".
+The user's intent will be passed to you between "==== BEGIN USER REQUEST ====" and "==== END USER REQUEST ====".
+Your job is to edit the cell based on the contents of the Srcbook and the user's intent.
+Act as a javascript expert coder, writing the best possible code you can. Focus on being elegant, concise, and clear.
+ONLY RETURN THE CODE, NO PREAMBULE, NO BACKTICKS, NO MARKDOWN, NO SUFFIX, ONLY THE JAVASCRIPT CODE.
+===== END FINAL INSTRUCTIONS ===

packages/api/prompts/code-updater-typescript.txt

@@ -0,0 +1,161 @@
+===== BEGIN INSTRUCTIONS CONTEXT =====
+
+You are tasked with editing a code snippet (or "cell") in a Srcbook."
+
+A Srcbook is a TypeScript notebook following a markdown-compatible format.
+
+## Srcbook spec
+
+Structure of a Srcbook:
+0. The language comment: `<!-- srcbook:{"language":"typescript"} -->`
+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. TypeScript code cells, which have a filename and source content.
+
+The user is already working on an existing Srcbook, and is asking you to edit a specific code cell. 
+The Srcbook contents will be passed to you as context, as well as the user request about what the edit intent they have for the code cell.
+===== END INSTRUCTIONS CONTEXT ======
+
+===== BEGIN EXAMPLE SRCBOOK =====
+<!-- srcbook:{"language":"typescript"} -->
+
+# Breadth-First Search (BFS) in TypeScript
+
+###### package.json
+
+```json
+{
+  "type": "module",
+  "dependencies": {},
+  "devDependencies": {
+    "tsx": "latest",
+    "typescript": "latest",
+    "@types/node": "^20.14.7"
+  }
+}
+```
+
+## Introduction to Breadth-First Search (BFS)
+
+Breadth-First Search (BFS) is an algorithm for traversing or searching tree or graph data structures. It starts at the tree root (or an arbitrary node of a graph) and explores the neighbor nodes at the present depth prior to moving on to nodes at the next depth level.
+
+BFS is particularly useful for finding the shortest path on unweighted graphs, and it can be implemented using a queue data structure.
+
+## BFS Algorithm Steps
+
+1. **Initialize**: Start from the root node and add it to the queue.
+2. **Dequeue**: Remove the front node from the queue and mark it as visited.
+3. **Enqueue**: Add all unvisited neighbors of the dequeued node to the queue.
+4. **Repeat**: Continue the process until the queue is empty.
+
+## BFS Implementation in TypeScript
+
+Let's implement BFS in TypeScript. We'll start by defining a simple graph structure and then implement the BFS algorithm.
+
+###### graph.ts
+
+```typescript
+export class Graph {
+  private adjacencyList: Map<number, number[]>;
+
+  constructor() {
+    this.adjacencyList = new Map();
+  }
+
+  addVertex(vertex: number) {
+    if (!this.adjacencyList.has(vertex)) {
+      this.adjacencyList.set(vertex, []);
+    }
+  }
+
+  addEdge(vertex1: number, vertex2: number) {
+    if (this.adjacencyList.has(vertex1) && this.adjacencyList.has(vertex2)) {
+      this.adjacencyList.get(vertex1)?.push(vertex2);
+      this.adjacencyList.get(vertex2)?.push(vertex1);
+    }
+  }
+
+  getNeighbors(vertex: number): number[] {
+    return this.adjacencyList.get(vertex) || [];
+  }
+}
+```
+
+In the above code, we define a `Graph` class with methods to add vertices and edges, and to retrieve the neighbors of a vertex.
+
+Next, let's implement the BFS algorithm.
+
+###### bfs.ts
+
+```typescript
+import { Graph } from './graph.ts';
+
+export function bfs(graph: Graph, startVertex: number): number[] {
+  const visited: Set<number> = new Set();
+  const queue: number[] = [startVertex];
+  const result: number[] = [];
+
+  while (queue.length > 0) {
+    const vertex = queue.shift()!;
+    if (!visited.has(vertex)) {
+      visited.add(vertex);
+      result.push(vertex);
+
+      const neighbors = graph.getNeighbors(vertex);
+      for (const neighbor of neighbors) {
+        if (!visited.has(neighbor)) {
+          queue.push(neighbor);
+        }
+      }
+    }
+  }
+
+  return result;
+}
+```
+
+In the `bfs` function, we use a queue to keep track of the vertices to be explored and a set to keep track of visited vertices. The function returns the order in which the vertices are visited.
+
+## Example Usage
+
+Let's create a graph and perform BFS on it.
+
+###### example.ts
+
+```typescript
+import { Graph } from './graph.ts';
+import { bfs } from './bfs.ts';
+
+const graph = new Graph();
+graph.addVertex(1);
+graph.addVertex(2);
+graph.addVertex(3);
+graph.addVertex(4);
+graph.addVertex(5);
+
+graph.addEdge(1, 2);
+graph.addEdge(1, 3);
+graph.addEdge(2, 4);
+graph.addEdge(3, 5);
+
+const bfsResult = bfs(graph, 1);
+console.log('BFS Traversal Order:', bfsResult);
+```
+
+In this example, we create a graph with 5 vertices and add edges between them. We then perform BFS starting from vertex 1 and print the traversal order.
+
+## Conclusion
+
+Breadth-First Search (BFS) is a fundamental algorithm for graph traversal. It is widely used in various applications, including finding the shortest path in unweighted graphs. In this srcbook, we implemented BFS in TypeScript and demonstrated its usage with a simple example.
+===== END EXAMPLE SRCBOOK =====
+
+===== BEGIN FINAL INSTRUCTIONS =====
+The user's Srcbook will be passed to you, surrounded with "==== BEGIN SRCBOOK ====" and "==== END SRCBOOK ====".
+The specific code cell they want updated will also be passed to you, surrounded with "==== BEGIN CODE CELL ====" and "==== END CODE CELL ====".
+The user's intent will be passed to you between "==== BEGIN USER REQUEST ====" and "==== END USER REQUEST ====".
+Your job is to edit the cell based on the contents of the Srcbook and the user's intent.
+Act as a TypeScript expert coder, writing the best possible code you can. Focus on being elegant, concise, and clear.
+ONLY RETURN THE CODE, NO PREAMBULE, NO BACKTICKS, NO MARKDOWN, NO SUFFIX, ONLY THE TYPESCRIPT CODE.
+===== END FINAL INSTRUCTIONS ===

packages/api/server/http.mts

@@ -13,7 +13,7 @@ import {
   sessionToResponse,
   listSessions,
 } from '../session.mjs';
-import { generateCell, generateSrcbook } from '../ai/srcbook-generator.mjs';
+import { generateCells, generateSrcbook } from '../ai/generate.mjs';
 import { disk } from '../utils.mjs';
 import { getConfig, updateConfig, getSecrets, addSecret, removeSecret } from '../config.mjs';
 import {
@@ -138,7 +138,7 @@ router.post('/sessions/:id/generate_cells', cors(), async (req, res) => {
   try {
     posthog.capture({ event: 'user generated cell with AI', properties: { query } });
     const session = await findSession(req.params.id);
-    const { error, errors, cells } = await generateCell(query, session, insertIdx);
+    const { error, errors, cells } = await generateCells(query, session, insertIdx);
     const result = error ? errors : cells;
     return res.json({ error, result });
   } catch (e) {