Generate ai fixes & improvements (#139)

* Support generating multiple cells

* Support multiple cells

* Update prompts for multi-cell generation

packages/api/ai/srcbook-generator.mts

@@ -80,7 +80,7 @@ export async function generateSrcbook(query: string): Promise<NoToolsGenerateTex
 type GenerateCellResult = {
   error: boolean;
   errors?: string[];
-  cell?: CellType;
+  cells?: CellType[];
 };
 export async function generateCell(
   query: string,
@@ -112,11 +112,6 @@ export async function generateCell(
   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] };
-    }
+    return { error: false, cells: decodeResult.cells };
   }
 }

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

@@ -1,8 +1,8 @@
 ===== BEGIN INSTRUCTIONS CONTEXT =====
 
-You are tasked with generating a Srcbook code cell for the user.
+You are tasked with generating Srcbook cells 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. 
+A Srcbook is a JavaScript notebook following a markdown-compatible format called .srcmd.
 
 ## Srcbook spec
 
@@ -14,8 +14,8 @@ Structure of a Srcbook:
   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.
+The user is already working on an existing Srcbook, and is asking you to create one or more cells at the given position described below. Cells can be code and markdown. If unspecified lean towards code rather than markdown.
+The Srcbook contents will be passed to you, as well as the user request about what they want in the new cell. 
 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.
@@ -41,7 +41,7 @@ Markdown cells are regular markdown. Just avoid using heading1 and heading6, as
 
 ## 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.
+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:
@@ -103,9 +103,10 @@ const token = auth(API_KEY);
 
 ===== 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.
+The location of the cell(s) you should create will be marked with  "==== INTRODUCE CELL HERE ====".
+The user's intent will be passed to you between "==== BEGIN USER REQUEST ====" and "==== END USER REQUEST ====".
+Your job is to write one or more cells. For code cells, the filename and the JavaScript code for this cell according to the Srcbook spec. 
+Lean towards code cells 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:
@@ -124,5 +125,5 @@ ws.on('open', () => {
 ```
 </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.
+ Act as a JavaScript expert, writing the best possible code you can. Focus on being elegant, concise, and clear.
 ===== END FINAL INSTRUCTIONS ===

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

@@ -1,8 +1,8 @@
 ===== BEGIN INSTRUCTIONS CONTEXT =====
 
-You are tasked with generating a Srcbook code cell for the user.
+You are tasked with generating Srcbook cells for the user.
 
-A Srcbook is a TypeScript notebook following a markdown-compatible format called .srcmd. It's an interactive and rich way of programming that follows the literate programming idea. 
+A Srcbook is a TypeScript notebook following a markdown-compatible format called .srcmd.
 
 ## Srcbook spec
 
@@ -14,8 +14,8 @@ Structure of a Srcbook:
   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 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.
+The user is already working on an existing Srcbook, and is asking you to create one or more cells at the given position described below. Cells can be code and markdown. If unspecified lean towards code rather than markdown.
+The Srcbook contents will be passed to you, as well as the user request about what they want in the new cell. 
 Each code cell needs to have a unique filename, as it maps to a file on disk.
 
 Code cells are valid TypeScript 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 typescript. Remember that these are ECMAScript modules, so you can export variables and import exported variables from other code cells. For example:
@@ -32,7 +32,7 @@ console.log(`Is string: ${B}`); // Output: Is string: No
 ```
 </example>
 
-Markdown cells are regular markdown. Just avoid using heading1 and heading6, as those are reserved by the Srcbook spec.
+Markdown cells are regular markdown. Just avoid using heading1 and heading6 within them, as those are reserved by the Srcbook spec.
 ===== END INSTRUCTIONS CONTEXT ======
 
 ===== BEGIN EXAMPLE SRCBOOK =====
@@ -166,14 +166,14 @@ In this example, we create a graph with 5 vertices and add edges between them. W
 ## 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 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 TypeScript code for this cell according to the Srcbook spec, or a Markdown cell. Lean towards a code cell if the user request is unclear.
+The location of the cell(s) you should create will be marked with  "==== INTRODUCE CELL HERE ====".
+The user's intent will be passed to you between "==== BEGIN USER REQUEST ====" and "==== END USER REQUEST ====".
+Your job is to write one or more cells. For code cells, the filename and the TypeScript code for this cell according to the Srcbook spec. 
+Lean towards code cells 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:
@@ -190,5 +190,5 @@ console.log(`Is string: ${B}`); // Output: Is string: No
 ```
 </example>
 
- Write the best possible code you can, as if you were an expert TypeScript engineer. Focus on being elegant, concise, clear. Keep things simple and explicit.
+ Act as a TypeScript expert coder, writing the best possible code you can. Focus on being elegant, concise, and clear.
 ===== END FINAL INSTRUCTIONS ===

packages/api/server/http.mts

@@ -129,16 +129,16 @@ router.post('/generate', cors(), async (req, res) => {
 });
 
 // Generate a cell using AI from a query string
-router.options('/sessions/:id/generate_cell', cors());
-router.post('/sessions/:id/generate_cell', cors(), async (req, res) => {
+router.options('/sessions/:id/generate_cells', cors());
+router.post('/sessions/:id/generate_cells', cors(), async (req, res) => {
   // @TODO: zod
   const { insertIdx, query } = req.body;
 
   try {
     posthog.capture({ event: 'user generated cell with AI', properties: { query } });
     const session = await findSession(req.params.id);
-    const { error, errors, cell } = await generateCell(query, session, insertIdx);
-    const result = error ? errors : cell;
+    const { error, errors, cells } = await generateCell(query, session, insertIdx);
+    const result = error ? errors : cells;
     return res.json({ error, result });
   } catch (e) {
     const error = e as unknown as Error;

packages/web/src/components/cells/generate-ai.tsx

@@ -1,26 +1,28 @@
 import { useState } from 'react';
+import { useNavigate } from 'react-router-dom';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { type CodeCellType, type MarkdownCellType } from '@srcbook/shared';
-import { generateCell } from '@/lib/server';
+import { generateCells } from '@/lib/server';
 import { CircleAlert, Trash2, Sparkles } from 'lucide-react';
 import { GenerateAICellType, SessionType } from '@/types';
 import { useCells } from '@/components/use-cell';
 import { cn } from '@/lib/utils';
 import { Button } from '@/components/ui/button';
-import DeleteCellWithConfirmation from '@/components/delete-cell-dialog';
 
 export default function GenerateAiCell(props: {
   cell: GenerateAICellType;
   insertIdx: number;
   session: SessionType;
-  onSuccess: (idx: number, cell: CodeCellType | MarkdownCellType) => void;
+  onSuccess: (idx: number, cells: Array<CodeCellType | MarkdownCellType>) => void;
+  hasOpenaiKey: boolean;
 }) {
-  const { cell, insertIdx, session, onSuccess } = props;
+  const { cell, insertIdx, session, onSuccess, hasOpenaiKey } = props;
   const [state, setState] = useState<'idle' | 'loading'>('idle');
   const { removeCell } = useCells();
   const [prompt, setPrompt] = useState('');
   const [error, setError] = useState<string | null>(null);
 
+  const navigate = useNavigate();
   useHotkeys(
     'mod+enter',
     () => {
@@ -33,7 +35,7 @@ export default function GenerateAiCell(props: {
   const generate = async () => {
     setError(null);
     setState('loading');
-    const { result, error } = await generateCell(session.id, {
+    const { result, error } = await generateCells(session.id, {
       query: prompt,
       insertIdx: insertIdx,
     });
@@ -60,30 +62,23 @@ export default function GenerateAiCell(props: {
           'ring-1 ring-sb-red-30 border-sb-red-30 hover:border-sb-red-30 focus-within:border-sb-red-30 focus-within:ring-sb-red-30',
       )}
     >
-      {error && (
-        <div className="flex items-center gap-2 absolute bottom-1 right-1 px-2.5 py-2 text-sb-red-80 bg-sb-red-30 rounded-sm">
-          <CircleAlert size={16} />
-          <p className="text-xs">{error}</p>
-        </div>
-      )}
       <div className="flex flex-col">
         <div className="p-1 w-full flex items-center justify-between z-10">
           <div className="flex items-center gap-2">
             <h5 className="pl-4 text-sm font-mono font-bold">Generate with AI</h5>
-            <DeleteCellWithConfirmation onDeleteCell={() => removeCell(cell)}>
-              <Button
-                variant="secondary"
-                size="icon"
-                className="border-secondary hover:border-muted"
-              >
-                <Trash2 size={16} />
-              </Button>
-            </DeleteCellWithConfirmation>
+            <Button
+              variant="secondary"
+              size="icon"
+              className="border-secondary hover:border-muted"
+              onClick={() => removeCell(cell)}
+            >
+              <Trash2 size={16} />
+            </Button>
           </div>
 
           <div>
             <Button
-              disabled={!prompt}
+              disabled={!prompt || !hasOpenaiKey}
               onClick={generate}
               variant={state === 'idle' ? 'default' : 'run'}
             >
@@ -92,17 +87,36 @@ export default function GenerateAiCell(props: {
           </div>
         </div>
 
-        <div className="flex items-start">
+        <div className={cn('flex items-start', error && 'border-b')}>
           <Sparkles size={16} className="m-2.5" />
 
           <textarea
             value={prompt}
             autoFocus
             onChange={(e) => setPrompt(e.target.value)}
             placeholder="Write a prompt..."
-            className="flex min-h-[60px] w-full rounded-sm px-3 py-2 text-sm placeholder:text-muted-foreground focus-visible:outline-none pl-0"
+            className="flex min-h-[80px] bg-transparent w-full rounded-sm px-3 py-2 text-sm placeholder:text-muted-foreground focus-visible:outline-none pl-0"
           />
         </div>
+
+        {error && (
+          <div className="flex items-center gap-2 m-2 px-2.5 py-2 text-sb-red-80 bg-sb-red-30 rounded-sm justify-center">
+            <CircleAlert size={16} />
+            <p className="text-xs line-clamp-1 ">{error}</p>
+          </div>
+        )}
+
+        {!hasOpenaiKey && (
+          <div className="flex items-center justify-between bg-sb-yellow-20 text-sb-yellow-80 rounded-sm text-sm p-1 m-3">
+            <p className="px-2">API key required</p>
+            <button
+              className="border border-sb-yellow-70 rounded-sm px-2 py-1 hover:border-sb-yellow-80 animate-all"
+              onClick={() => navigate('/settings')}
+            >
+              Settings
+            </button>
+          </div>
+        )}
       </div>
     </div>
   );

packages/web/src/lib/server.ts

@@ -1,4 +1,4 @@
-import { CodeLanguageType, type CodeCellType } from '@srcbook/shared';
+import type { CodeLanguageType, MarkdownCellType, CodeCellType } from '@srcbook/shared';
 import { SessionType, FsObjectResultType, ExampleSrcbookType } from '@/types';
 
 const API_BASE_URL = 'http://localhost:2150/api';
@@ -116,15 +116,15 @@ export async function generateSrcbook(
   return response.json();
 }
 
-type GenerateCellRequestType = { insertIdx: number; query: string };
-type GenerateCellResponseType =
+type GenerateCellsRequestType = { insertIdx: number; query: string };
+type GenerateCellsResponseType =
   | { error: true; result: string }
-  | { error: false; result: CodeCellType };
-export async function generateCell(
+  | { error: false; result: Array<CodeCellType | MarkdownCellType> };
+export async function generateCells(
   sessionId: string,
-  request: GenerateCellRequestType,
-): Promise<GenerateCellResponseType> {
-  const response = await fetch(API_BASE_URL + '/sessions/' + sessionId + '/generate_cell', {
+  request: GenerateCellsRequestType,
+): Promise<GenerateCellsResponseType> {
+  const response = await fetch(API_BASE_URL + '/sessions/' + sessionId + '/generate_cells', {
     method: 'POST',
     headers: { 'content-type': 'application/json' },
     body: JSON.stringify(request),

packages/web/src/routes/session.tsx

@@ -12,8 +12,8 @@ import {
   TitleCellType,
   PackageJsonCellType,
 } from '@srcbook/shared';
-import { loadSession } from '@/lib/server';
-import { SessionType, type GenerateAICellType } from '@/types';
+import { loadSession, getConfig } from '@/lib/server';
+import type { SessionType, GenerateAICellType, SettingsType } from '@/types';
 import TitleCell from '@/components/cells/title';
 import MarkdownCell from '@/components/cells/markdown';
 import GenerateAiCell from '@/components/cells/generate-ai';
@@ -27,12 +27,16 @@ import useEffectOnce from '@/components/use-effect-once';
 import { cn } from '@/lib/utils';
 
 async function loader({ params }: LoaderFunctionArgs) {
-  const { result: session } = await loadSession({ id: params.id! });
-  return { session };
+  const [{ result: config }, { result: session }] = await Promise.all([
+    getConfig(),
+    loadSession({ id: params.id! }),
+  ]);
+
+  return { config, session };
 }
 
 function SessionPage() {
-  const { session } = useLoaderData() as { session: SessionType };
+  const { session, config } = useLoaderData() as { session: SessionType; config: SettingsType };
 
   const channelRef = useRef(SessionChannel.create(session.id));
   const channel = channelRef.current;
@@ -58,12 +62,12 @@ function SessionPage() {
 
   return (
     <CellsProvider initialCells={session.cells}>
-      <Session session={session} channel={channelRef.current} />
+      <Session session={session} channel={channelRef.current} config={config} />
     </CellsProvider>
   );
 }
 
-function Session(props: { session: SessionType; channel: SessionChannel }) {
+function Session(props: { session: SessionType; channel: SessionChannel; config: SettingsType }) {
   const { session, channel } = props;
 
   const {
@@ -162,17 +166,21 @@ function Session(props: { session: SessionType; channel: SessionChannel }) {
     }
   }
 
-  async function insertGeneratedCell(idx: number, cell: CodeCellType | MarkdownCellType) {
-    let newCell;
-    switch (cell.type) {
-      case 'code':
-        newCell = createCodeCell(idx, session.metadata.language, cell);
-        break;
-      case 'markdown':
-        newCell = createMarkdownCell(idx, cell);
-        break;
+  async function insertGeneratedCells(idx: number, cells: Array<CodeCellType | MarkdownCellType>) {
+    for (let i = 0; i < cells.length; i++) {
+      const cell = cells[i];
+      const insertIdx = idx + i;
+      let newCell;
+      switch (cell.type) {
+        case 'code':
+          newCell = createCodeCell(insertIdx, session.metadata.language, cell);
+          break;
+        case 'markdown':
+          newCell = createMarkdownCell(insertIdx, cell);
+          break;
+      }
+      channel.push('cell:create', { sessionId: session.id, index: insertIdx, cell: newCell });
     }
-    channel.push('cell:create', { sessionId: session.id, index: idx, cell: newCell });
   }
 
   // TOOD: We need to stop treating titles and package.json as cells.
@@ -222,7 +230,8 @@ function Session(props: { session: SessionType; channel: SessionChannel }) {
                 cell={cell}
                 session={session}
                 insertIdx={idx + 2}
-                onSuccess={insertGeneratedCell}
+                onSuccess={insertGeneratedCells}
+                hasOpenaiKey={!!props.config.openaiKey}
               />
             )}
           </div>
@@ -280,7 +289,7 @@ function InsertCellDivider(props: {
             className="border-none rounded-md rounded-l-none"
             onClick={props.createGenerateAiCodeCell}
           >
-            Create with AI
+            Generate with AI
           </Button>
         </div>
       </div>