@marimo-team/islands 0.17.0 → 0.17.1

package/src/core/ai/tools/edit-notebook-tool.ts

@@ -0,0 +1,239 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { EditorView } from "@codemirror/view";
+import { z } from "zod";
+import { scrollAndHighlightCell } from "@/components/editor/links/cell-link";
+import {
+  type CellPosition as NotebookCellPosition,
+  type NotebookState,
+  notebookAtom,
+} from "@/core/cells/cells";
+import { CellId } from "@/core/cells/ids";
+import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
+import type { CellColumnId } from "@/utils/id-tree";
+import { stagedAICellsAtom } from "../staged-cells";
+import {
+  type AiTool,
+  type ToolDescription,
+  ToolExecutionError,
+  type ToolNotebookContext,
+  type ToolOutputBase,
+  toolOutputBaseSchema,
+} from "./base";
+import type { CopilotMode } from "./registry";
+
+const description: ToolDescription = {
+  baseDescription:
+    "Perform editing operations on the current notebook. You should prefer to create new cells unless you need to edit existing cells. Call this tool multiple times to perform multiple edits. Separate code into logical individual cells to take advantage of the notebook's reactive execution model.",
+  prerequisites: [
+    "If you are updating existing cells, you need the cellIds. If they are not known, call the lightweight_cell_map_tool to find out.",
+  ],
+  additionalInfo: `
+  Args:
+    edit (object): The editing operation to perform. Must be one of:
+    - update_cell: Update the code of an existing cell, pass CellId and the new code.
+    - add_cell: Add a new cell to the notebook. The position of the new cell is specified by the position argument.
+        Pass "end" to add the new cell at the end of the notebook. 
+        Pass { cellId: cellId, before: true } to add the new cell before the specified cell. And before: false if after the specified cell.
+        Pass { type: "end", columnIndex: number } to add the new cell at the end of a specified column index. The column index is 0-based.
+    - delete_cell: Delete an existing cell, pass CellId. For deleting cells, the user needs to accept the deletion to actually delete the cell, so you may still see the cell in the notebook on subsequent edits which is fine.
+
+    For adding code, use the following guidelines:
+    - Markdown cells: use mo.md(f"""{content}""") function to insert content.
+    - SQL cells: use mo.sql(f"""{content}""") function to insert content. If a database engine is specified, use mo.sql(f"""{content}""", engine=engine) instead.
+
+    Returns:
+    - A result object containing standard tool metadata.`,
+};
+
+type CellPosition =
+  | { cellId: CellId; before: boolean }
+  | { type: "end"; columnIndex: number }
+  | "end";
+
+const editNotebookSchema = z.object({
+  edit: z.discriminatedUnion("type", [
+    z.object({
+      type: z.literal("update_cell"),
+      cellId: z.string() as unknown as z.ZodType<CellId>,
+      code: z.string(),
+    }),
+    z.object({
+      type: z.literal("add_cell"),
+      position: z.union([
+        z.object({
+          cellId: z.string() as unknown as z.ZodType<CellId>,
+          before: z.boolean(),
+        }),
+        z.object({
+          type: z.literal("end"),
+          columnIndex: z.number(),
+        }),
+        z.literal("end"),
+      ]) satisfies z.ZodType<CellPosition>,
+      code: z.string(),
+    }),
+    z.object({
+      type: z.literal("delete_cell"),
+      cellId: z.string() as unknown as z.ZodType<CellId>,
+    }),
+  ]),
+});
+
+type EditNotebookInput = z.infer<typeof editNotebookSchema>;
+type EditOperation = EditNotebookInput["edit"];
+export type EditType = EditOperation["type"];
+
+export class EditNotebookTool
+  implements AiTool<EditNotebookInput, ToolOutputBase>
+{
+  readonly name = "edit_notebook_tool";
+  readonly description = description;
+  readonly schema = editNotebookSchema;
+  readonly outputSchema = toolOutputBaseSchema;
+  readonly mode: CopilotMode[] = ["agent"];
+
+  handler = async (
+    { edit }: EditNotebookInput,
+    toolContext: ToolNotebookContext,
+  ): Promise<ToolOutputBase> => {
+    const { addStagedCell, createNewCell, store } = toolContext;
+
+    switch (edit.type) {
+      case "update_cell": {
+        const { cellId, code } = edit;
+
+        const notebook = store.get(notebookAtom);
+        this.validateCellIdExists(cellId, notebook);
+        const editorView = this.getCellEditorView(cellId, notebook);
+
+        scrollAndHighlightCell(cellId);
+
+        // If previous code exists, we don't want to replace it, it means there is a new edit on top of the previous edit
+        // Keep the original code
+        const stagedCell = store.get(stagedAICellsAtom).get(cellId);
+        const currentCellCode = editorView.state.doc.toString();
+        const previousCode =
+          stagedCell?.type === "update_cell" ||
+          stagedCell?.type === "delete_cell"
+            ? stagedCell.previousCode
+            : currentCellCode;
+
+        addStagedCell({
+          cellId,
+          edit: { type: "update_cell", previousCode: previousCode },
+        });
+
+        updateEditorCodeFromPython(editorView, code);
+
+        break;
+      }
+      case "add_cell": {
+        const { position, code } = edit;
+
+        // By default, add the new cell to the end of the notebook
+        let notebookPosition: NotebookCellPosition = "__end__";
+        let before = false;
+        const newCellId = CellId.create();
+
+        if (typeof position === "object") {
+          const notebook = store.get(notebookAtom);
+          if ("cellId" in position) {
+            this.validateCellIdExists(position.cellId, notebook);
+            notebookPosition = position.cellId;
+            before = position.before;
+          } else if ("columnIndex" in position) {
+            const columnId = this.getColumnId(position.columnIndex, notebook);
+            notebookPosition = { type: "__end__", columnId };
+          }
+        }
+
+        createNewCell({
+          cellId: notebookPosition,
+          before,
+          code,
+          newCellId,
+        });
+
+        addStagedCell({
+          cellId: newCellId,
+          edit: { type: "add_cell" },
+        });
+
+        // Scroll into view
+        scrollAndHighlightCell(newCellId);
+
+        break;
+      }
+      case "delete_cell": {
+        const { cellId } = edit;
+
+        const notebook = store.get(notebookAtom);
+        this.validateCellIdExists(cellId, notebook);
+
+        const editorView = this.getCellEditorView(cellId, notebook);
+        const currentCellCode = editorView.state.doc.toString();
+
+        // Add to staged AICells - don't actually delete the cell yet
+        addStagedCell({
+          cellId,
+          edit: { type: "delete_cell", previousCode: currentCellCode },
+        });
+
+        scrollAndHighlightCell(cellId);
+        break;
+      }
+    }
+    return {
+      status: "success",
+      next_steps: ["If you need to perform more edits, call this tool again."],
+    };
+  };
+
+  private validateCellIdExists(cellId: CellId, notebook: NotebookState) {
+    const cellIds = notebook.cellIds;
+    if (!cellIds.getColumns().some((column) => column.idSet.has(cellId))) {
+      throw new ToolExecutionError(
+        "Cell not found",
+        "CELL_NOT_FOUND",
+        false,
+        "Check which cells exist in the notebook",
+      );
+    }
+  }
+
+  private getColumnId(
+    columnIndex: number,
+    notebook: NotebookState,
+  ): CellColumnId {
+    const cellIds = notebook.cellIds;
+    const columns = cellIds.getColumns();
+
+    if (columnIndex < 0 || columnIndex >= columns.length) {
+      throw new ToolExecutionError(
+        "Column index is out of range",
+        "COLUMN_INDEX_OUT_OF_RANGE",
+        true,
+        "Choose a column index between 0 and the number of columns in the notebook (0-based)",
+      );
+    }
+    return columns[columnIndex].id;
+  }
+
+  private getCellEditorView(
+    cellId: CellId,
+    notebook: NotebookState,
+  ): EditorView {
+    const cellHandles = notebook.cellHandles;
+    const cellHandle = cellHandles[cellId].current;
+    if (!cellHandle?.editorView) {
+      throw new ToolExecutionError(
+        "Cell editor not found",
+        "CELL_EDITOR_NOT_FOUND",
+        false,
+        "Internal error, ask the user to report this error",
+      );
+    }
+    return cellHandle.editorView;
+  }
+}

package/src/core/ai/tools/registry.ts

@@ -3,8 +3,14 @@
 import type { components } from "@marimo-team/marimo-api";
 import { Memoize } from "typescript-memoize";
 import { type ZodObject, z } from "zod";
-import { type AiTool, ToolExecutionError } from "./base";
-import { TestFrontendTool } from "./sample-tool";
+import {
+  type AiTool,
+  ToolExecutionError,
+  type ToolNotebookContext,
+} from "./base";
+import { EditNotebookTool } from "./edit-notebook-tool";
+import { RunStaleCellsTool } from "./run-cells-tool";
+import { formatToolDescription } from "./utils";
 
 export type AnyZodObject = ZodObject<z.ZodRawShape>;
 
@@ -49,6 +55,7 @@ export class FrontendToolRegistry {
   async invoke<TName extends string>(
     toolName: TName,
     rawArgs: unknown,
+    toolContext: ToolNotebookContext,
   ): Promise<InvokeResult<TName>> {
     const tool = this.getToolOrThrow(toolName);
     const handler = tool.handler;
@@ -65,7 +72,7 @@ export class FrontendToolRegistry {
       const args = inputResponse.data;
 
       // Call the handler
-      const rawOutput = await handler(args);
+      const rawOutput = await handler(args, toolContext);
 
       // Parse output
       const response = await outputSchema.safeParseAsync(rawOutput);
@@ -114,7 +121,7 @@ export class FrontendToolRegistry {
   getToolSchemas(): FrontendToolDefinition[] {
     return [...this.tools.values()].map((tool) => ({
       name: tool.name,
-      description: tool.description,
+      description: formatToolDescription(tool.description),
       parameters: z.toJSONSchema(tool.schema),
       source: "frontend",
       mode: tool.mode,
@@ -123,6 +130,6 @@ export class FrontendToolRegistry {
 }
 
 export const FRONTEND_TOOL_REGISTRY = new FrontendToolRegistry([
-  ...(import.meta.env.DEV ? [new TestFrontendTool()] : []),
-  // ADD MORE TOOLS HERE
+  new EditNotebookTool(),
+  new RunStaleCellsTool(),
 ]);

package/src/core/ai/tools/run-cells-tool.ts

@@ -0,0 +1,233 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { FileUIPart } from "ai";
+import { z } from "zod";
+import { runCells } from "@/components/editor/cell/useRunCells";
+import { notebookAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import { staleCellIds } from "@/core/cells/utils";
+import { type JotaiStore, waitFor } from "@/core/state/jotai";
+import {
+  type BaseOutput,
+  getCellContextData,
+} from "../context/providers/cell-output";
+import {
+  type AiTool,
+  type EmptyToolInput,
+  type ToolDescription,
+  type ToolNotebookContext,
+  type ToolOutputBase,
+  toolOutputBaseSchema,
+} from "./base";
+import type { CopilotMode } from "./registry";
+
+interface CellOutput {
+  consoleOutput?: string;
+  // consoleAttachments?: FileUIPart[];
+  cellOutput?: string;
+  // cellAttachments?: FileUIPart[];
+}
+
+// Must use Record instead of Map because Map serializes to JSON as {}
+interface RunStaleCellsOutput extends ToolOutputBase {
+  cellsToOutput?: Record<CellId, CellOutput | null>;
+}
+
+const description: ToolDescription = {
+  baseDescription:
+    "Run cells in the current notebook that are stale. Stale cells are cells that have been edited or newly added. You can run this tool after editing the notebook or when requested. The output of the ran cells will be returned alongside metadata.",
+};
+
+const filePartSchema = z.object({
+  type: z.literal("file"),
+  url: z.string(),
+  mediaType: z.string(),
+}) satisfies z.ZodType<FileUIPart>;
+
+export class RunStaleCellsTool
+  implements AiTool<EmptyToolInput, RunStaleCellsOutput>
+{
+  readonly name = "run_stale_cells_tool";
+  readonly description = description;
+  readonly schema = z.object({});
+  readonly outputSchema = toolOutputBaseSchema.extend({
+    cellsToOutput: z
+      .record(
+        z.string(),
+        z
+          .object({
+            consoleOutput: z.string().optional(),
+            cellOutput: z.string().optional(),
+            consoleAttachments: z.array(filePartSchema).optional(),
+            cellAttachments: z.array(filePartSchema).optional(),
+          })
+          .nullable(),
+      )
+      .optional(),
+  }) satisfies z.ZodType<RunStaleCellsOutput>;
+  readonly mode: CopilotMode[] = ["agent"];
+
+  handler = async (
+    _args: EmptyToolInput,
+    toolContext: ToolNotebookContext,
+  ): Promise<RunStaleCellsOutput> => {
+    const { prepareForRun, sendRun, store } = toolContext;
+
+    const notebook = store.get(notebookAtom);
+    const staleCells = staleCellIds(notebook);
+
+    if (staleCells.length === 0) {
+      return {
+        status: "success",
+        message: "No stale cells found.",
+      };
+    }
+
+    await runCells({
+      cellIds: staleCells,
+      sendRun: sendRun,
+      prepareForRun,
+      notebook,
+    });
+
+    // Wait for all cells to finish executing
+    const allCellsFinished = await this.waitForCellsToFinish(store, staleCells);
+    if (!allCellsFinished) {
+      return {
+        status: "success",
+        message:
+          "No output was returned because some cells have not finished executing",
+      };
+    }
+
+    // Get notebook state after cells have finished
+    const updatedNotebook = store.get(notebookAtom);
+
+    const cellsToOutput = new Map<CellId, CellOutput | null>();
+    let resultMessage = "";
+
+    for (const cellId of staleCells) {
+      const cellContextData = getCellContextData(cellId, updatedNotebook, {
+        includeConsoleOutput: true,
+      });
+
+      let cellOutputString: string | undefined;
+      // let cellAttachments: FileUIPart[] | undefined;
+      let consoleOutputString: string | undefined;
+      // let consoleAttachments: FileUIPart[] | undefined;
+
+      const cellOutput = cellContextData.cellOutput;
+      const consoleOutputs = cellContextData.consoleOutputs;
+      if (!cellOutput && !consoleOutputs) {
+        // Set null to show no output
+        cellsToOutput.set(cellId, null);
+        continue;
+      }
+
+      if (cellOutput) {
+        cellOutputString = this.formatOutputString(cellOutput);
+        // cellAttachments = await getAttachmentsForOutputs(
+        //   [cellOutput],
+        //   cellId,
+        //   cellContextData.cellName,
+        // );
+      }
+
+      if (consoleOutputs) {
+        // consoleAttachments = await getAttachmentsForOutputs(
+        //   consoleOutputs,
+        //   cellId,
+        //   cellContextData.cellName,
+        // );
+        consoleOutputString = consoleOutputs
+          .map((output) => this.formatOutputString(output))
+          .join("\n");
+        resultMessage +=
+          "Console output represents the stdout or stderr of the cell (eg. print statements).";
+      }
+
+      cellsToOutput.set(cellId, {
+        cellOutput: cellOutputString,
+        consoleOutput: consoleOutputString,
+        // cellAttachments: cellAttachments,
+        // consoleAttachments: consoleAttachments,
+      });
+    }
+
+    if (cellsToOutput.size === 0) {
+      return {
+        status: "success",
+        message: "Stale cells have been run. No output was returned.",
+      };
+    }
+
+    return {
+      status: "success",
+      cellsToOutput: Object.fromEntries(cellsToOutput),
+      message: resultMessage === "" ? undefined : resultMessage,
+      next_steps: [
+        "Review the output of the cells, if you need to make any changes, you can run this tool again, the CellId is the key of the result object if you want to edit the cell.",
+      ],
+    };
+  };
+
+  private formatOutputString(cellOutput: BaseOutput): string {
+    let outputString = "";
+    const { outputType, processedContent, imageUrl, output } = cellOutput;
+    if (outputType === "text") {
+      outputString += "Output:\n";
+      if (processedContent) {
+        outputString += processedContent;
+      } else if (typeof output.data === "string") {
+        outputString += output.data;
+      } else {
+        outputString += JSON.stringify(output.data);
+      }
+    } else if (outputType === "media") {
+      outputString += `Media Output: Contains ${output.mimetype} content`;
+      if (imageUrl) {
+        outputString += `\nImage URL: ${imageUrl}`;
+      }
+    }
+    return outputString;
+  }
+
+  /**
+   * Wait for cells to finish executing (status becomes "idle")
+   * Returns true if all cells finished executing, false if the timeout was reached
+   */
+  private async waitForCellsToFinish(
+    store: JotaiStore,
+    cellIds: CellId[],
+    timeout = 30_000,
+  ): Promise<boolean> {
+    const checkAllFinished = (
+      notebook: ReturnType<typeof notebookAtom.read>,
+    ) => {
+      return cellIds.every((cellId) => {
+        const cellRuntime = notebook.cellRuntime[cellId];
+        return (
+          cellRuntime?.status !== "running" && cellRuntime?.status !== "queued"
+        );
+      });
+    };
+
+    // If already finished, return immediately
+    if (checkAllFinished(store.get(notebookAtom))) {
+      return true;
+    }
+
+    // Wait for notebook state changes with timeout
+    try {
+      await Promise.race([
+        waitFor(notebookAtom, checkAllFinished),
+        new Promise<never>((_, reject) =>
+          setTimeout(() => reject(new Error("timeout")), timeout),
+        ),
+      ]);
+      return true;
+    } catch {
+      return false;
+    }
+  }
+}

package/src/core/ai/tools/sample-tool.ts

@@ -3,21 +3,23 @@
 import { z } from "zod";
 import {
   type AiTool,
+  type ToolDescription,
   ToolExecutionError,
   type ToolOutputBase,
   toolOutputBaseSchema,
 } from "./base";
 import type { CopilotMode } from "./registry";
 
-const description = `
-Test frontend tool that returns a greeting message.
+const description: ToolDescription = {
+  baseDescription: "Test frontend tool that returns a greeting message.",
+  additionalInfo: `
+  Args:
+    - name (string): The name to include in the greeting.
 
-Args:
-- name (string): The name to include in the greeting.
-
-Returns:
-- Output with data containing the greeting message.
-`;
+  Returns:
+    - Output with data containing the greeting message.
+  `,
+};
 
 interface Input {
   name: string;

package/src/core/ai/tools/utils.ts

@@ -0,0 +1,23 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { ToolDescription } from "./base";
+
+export function formatToolDescription(description: ToolDescription): string {
+  let result = description.baseDescription;
+  if (description.whenToUse) {
+    result += `\n\n## When to use:\n- ${description.whenToUse.join("\n- ")}`;
+  }
+  if (description.avoidIf) {
+    result += `\n\n## Avoid if:\n- ${description.avoidIf.join("\n- ")}`;
+  }
+  if (description.prerequisites) {
+    result += `\n\n## Prerequisites:\n- ${description.prerequisites.join("\n- ")}`;
+  }
+  if (description.sideEffects) {
+    result += `\n\n## Side effects:\n- ${description.sideEffects.join("\n- ")}`;
+  }
+  if (description.additionalInfo) {
+    result += `\n\n## Additional info:\n- ${description.additionalInfo}`;
+  }
+  return result;
+}

package/src/core/cells/__tests__/cells.test.ts

@@ -2109,7 +2109,7 @@ describe("cell reducer", () => {
     actions.addSetupCellIfDoesntExist({ code: "# Updated setup code" });
 
     // Check that the same setup cell was updated, not duplicated
-    expect(state.cellData[SETUP_CELL_ID].code).toBe("# Setup code");
+    expect(state.cellData[SETUP_CELL_ID].code).toBe("# Updated setup code");
     expect(state.cellData[SETUP_CELL_ID].edited).toBe(true);
     expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
   });
@@ -2143,6 +2143,24 @@ describe("cell reducer", () => {
     expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
   });
 
+  it("can delete and then create a new setup cell", () => {
+    // Create the setup cell
+    actions.addSetupCellIfDoesntExist({ code: "# Setup code" });
+
+    // Delete the setup cell
+    actions.deleteCell({ cellId: SETUP_CELL_ID });
+
+    // Create a new setup cell
+    actions.addSetupCellIfDoesntExist({ code: "# New code" });
+
+    // Check that the new setup cell was created
+    expect(state.cellData[SETUP_CELL_ID].id).toBe(SETUP_CELL_ID);
+    expect(state.cellData[SETUP_CELL_ID].name).toBe("setup");
+    expect(state.cellData[SETUP_CELL_ID].code).toBe("# New code");
+    expect(state.cellData[SETUP_CELL_ID].edited).toBe(true);
+    expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
+  });
+
   it("can clear all outputs", () => {
     // Add a cell and give it output
     actions.createNewCell({

package/src/core/cells/__tests__/session.test.ts

@@ -210,8 +210,18 @@ describe("notebookStateFromSession", () => {
 
     it("handles console outputs in session cell", () => {
       const consoleOutputs = [
-        { type: "stream", name: "stdout", text: "Hello stdout" } as const,
-        { type: "stream", name: "stderr", text: "Hello stderr" } as const,
+        {
+          type: "stream",
+          name: "stdout",
+          text: "Hello stdout",
+          mimetype: "text/plain",
+        } as const,
+        {
+          type: "stream",
+          name: "stderr",
+          text: "Hello stderr",
+          mimetype: "text/plain",
+        } as const,
       ];
       const session = createSession([
         createSessionCell("cell-1", [], consoleOutputs),
@@ -356,7 +366,14 @@ describe("notebookStateFromSession", () => {
         createSessionCell(
           "cell-1",
           [],
-          [{ type: "stream", name: "stdout", text: "output" }],
+          [
+            {
+              type: "stream",
+              name: "stdout",
+              text: "output",
+              mimetype: "text/plain",
+            },
+          ],
         ),
       ]);
       const notebook = createNotebook([

package/src/core/cells/cells.ts

@@ -145,13 +145,18 @@ export function initialNotebookState(): NotebookState {
   });
 }
 
+/** The target cell ID to create a new cell relative to. Can be:
+ * - A CellId string for an existing cell
+ * - "__end__" to append at the end of the first column
+ * - {type: "__end__", columnId} to append at the end of a specific column
+ */
+export type CellPosition =
+  | CellId
+  | "__end__"
+  | { type: "__end__"; columnId: CellColumnId };
+
 export interface CreateNewCellAction {
-  /** The target cell ID to create a new cell relative to. Can be:
-   * - A CellId string for an existing cell
-   * - "__end__" to append at the end of the first column
-   * - {type: "__end__", columnId} to append at the end of a specific column
-   */
-  cellId: CellId | "__end__" | { type: "__end__"; columnId: CellColumnId };
+  cellId: CellPosition;
   /** Whether to insert before (true) or after (false) the target cell */
   before: boolean;
   /** Initial code content for the new cell */
@@ -1331,7 +1336,7 @@ const {
     }
 
     // First check if setup cell already exists
-    if (SETUP_CELL_ID in state.cellData) {
+    if (SETUP_CELL_ID in state.cellIds) {
       // Just focus on the existing setup cell
       return {
         ...state,

package/src/core/cells/session.ts

@@ -229,7 +229,7 @@ function createCellRuntimeFromSession(
       return {
         channel: consoleOutput.name === "stderr" ? "stderr" : "stdout",
         data: consoleOutput.text,
-        mimetype: "text/plain",
+        mimetype: consoleOutput.mimetype ?? "text/plain",
         timestamp: DEFAULT_TIMESTAMP,
       };
     }),