@marimo-team/frontend 0.15.1-dev34 → 0.15.1-dev36

package/src/core/ai/context/providers/__tests__/error.test.ts

@@ -0,0 +1,328 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+import { createStore } from "jotai";
+import { beforeEach, describe, expect, it } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { notebookAtom } from "@/core/cells/cells";
+import { type CellId, CellId as CellIdClass } from "@/core/cells/ids";
+import { ErrorContextProvider } from "../error";
+
+describe("ErrorContextProvider", () => {
+  let provider: ErrorContextProvider;
+  let store: ReturnType<typeof createStore>;
+
+  beforeEach(() => {
+    store = createStore();
+    provider = new ErrorContextProvider(store);
+  });
+
+  const createMockNotebookWithErrors = (
+    errors: Parameters<typeof MockNotebook.notebookStateWithErrors>[0],
+  ) => {
+    const notebookState = MockNotebook.notebookStateWithErrors(errors);
+    store.set(notebookAtom, notebookState);
+  };
+
+  describe("provider properties", () => {
+    it("should have correct provider properties", () => {
+      expect(provider.title).toBe("Errors");
+      expect(provider.mentionPrefix).toBe("@");
+      expect(provider.contextType).toBe("error");
+    });
+  });
+
+  describe("getItems", () => {
+    it("should return empty array when no errors", () => {
+      const items = provider.getItems();
+      expect(items).toEqual([]);
+    });
+
+    it("should return error items when errors exist", () => {
+      const cellId1 = CellIdClass.create();
+      const cellId2 = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId: cellId1,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+        {
+          cellId: cellId2,
+          cellName: "Cell 2",
+          errorData: [MockNotebook.errors.exception("Runtime error")],
+        },
+      ]);
+
+      const items = provider.getItems();
+
+      expect(items).toHaveLength(1);
+      expect(items[0]).toMatchObject({
+        name: "Errors",
+        type: "error",
+        description: "All errors in the notebook",
+        data: {
+          type: "all-errors",
+          errors: [
+            {
+              cellId: cellId1,
+              cellName: "Cell 1",
+              errorData: [{ type: "syntax", msg: "Invalid syntax" }],
+            },
+            {
+              cellId: cellId2,
+              cellName: "Cell 2",
+              errorData: [{ type: "exception", msg: "Runtime error" }],
+            },
+          ],
+        },
+      });
+    });
+
+    it("should handle cells without names", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      expect(items[0].data.errors[0].cellName).toBe("");
+    });
+  });
+
+  describe("formatCompletion", () => {
+    it("should format completion for all-errors type", () => {
+      const cellId1 = CellIdClass.create();
+      const cellId2 = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId: cellId1,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+        {
+          cellId: cellId2,
+          cellName: "Cell 2",
+          errorData: [MockNotebook.errors.exception("Runtime error")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const completion = provider.formatCompletion(items[0]);
+
+      expect(completion).toMatchObject({
+        label: "@Errors",
+        displayLabel: "Errors",
+        detail: "2 errors",
+        boost: 2, // Boosts.ERROR
+        type: "error",
+        apply: "@Errors",
+        section: "Errors",
+      });
+
+      // Test the info function
+      expect(completion.info).toBeDefined();
+      if (typeof completion.info === "function") {
+        const infoResult = completion.info(completion);
+        if (
+          infoResult &&
+          typeof infoResult === "object" &&
+          "dom" in infoResult
+        ) {
+          const infoElement = infoResult.dom as HTMLElement;
+          expect(infoElement.tagName).toBe("DIV");
+          expect(infoElement.textContent).toContain("Errors");
+          expect(infoElement.textContent).toContain("2 errors");
+        } else if (infoResult && "tagName" in (infoResult as any)) {
+          const infoElement = infoResult as HTMLElement;
+          expect(infoElement.tagName).toBe("DIV");
+          expect(infoElement.textContent).toContain("Errors");
+          expect(infoElement.textContent).toContain("2 errors");
+        }
+      }
+    });
+
+    it("should handle single error correctly", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const completion = provider.formatCompletion(items[0]);
+      expect(completion.detail).toBe("1 error");
+    });
+
+    it("should handle fallback for unknown error types", () => {
+      const item = {
+        uri: "error://unknown",
+        name: "Unknown Error",
+        type: "error" as const,
+        data: {
+          type: "unknown" as any,
+          errors: [],
+        },
+        description: "Unknown error type",
+      };
+
+      const completion = provider.formatCompletion(item);
+      expect(completion).toMatchObject({
+        label: "Error",
+        displayLabel: "Error",
+        boost: 2, // Boosts.ERROR
+      });
+    });
+  });
+
+  describe("formatContext", () => {
+    it("should format context for basic errors", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[0]);
+      expect(context).toMatchSnapshot("basic-error-context");
+    });
+
+    it("should format context for multiple error types", () => {
+      const cellId1 = "cell-1" as CellId;
+      const cellId2 = "cell-2" as CellId;
+
+      createMockNotebookWithErrors([
+        {
+          cellId: cellId1,
+          cellName: "Cell 1",
+          errorData: [
+            MockNotebook.errors.syntax("Invalid syntax"),
+            MockNotebook.errors.exception("Runtime error"),
+          ],
+        },
+        {
+          cellId: cellId2,
+          cellName: "",
+          errorData: [
+            MockNotebook.errors.cycle(),
+            MockNotebook.errors.multipleDefs("variable_x"),
+          ],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[0]);
+
+      // Check for expected content instead of exact snapshot match due to random cell IDs
+      expect(context).toMatchInlineSnapshot(`
+        "<error name="Cell 1" description="Invalid syntax
+        Runtime error"></error>
+
+        <error name="Cell cell-2" description="This cell is in a cycle
+        The variable 'variable_x' was defined by another cell"></error>"
+      `);
+    });
+
+    it("should handle cells without names", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[0]);
+      expect(context).toContain(`Cell ${cellId}`);
+    });
+  });
+
+  describe("error description handling", () => {
+    it("should handle all error types correctly", () => {
+      // Note: ancestor-* error types are intentionally filtered out by cellErrorsAtom
+      const errorTypes = [
+        {
+          error: MockNotebook.errors.setupRefs(),
+          expected: "The setup cell cannot have references",
+        },
+        {
+          error: MockNotebook.errors.cycle(),
+          expected: "This cell is in a cycle",
+        },
+        {
+          error: MockNotebook.errors.multipleDefs("var_x"),
+          expected: "The variable 'var_x' was defined by another cell",
+        },
+        {
+          error: MockNotebook.errors.importStar("Import star error"),
+          expected: "Import star error",
+        },
+        {
+          error: MockNotebook.errors.exception("Runtime exception"),
+          expected: "Runtime exception",
+        },
+        {
+          error: MockNotebook.errors.strictException("Strict exception", "ref"),
+          expected: "Strict exception",
+        },
+        {
+          error: MockNotebook.errors.interruption(),
+          expected: "This cell was interrupted and needs to be re-run",
+        },
+        {
+          error: MockNotebook.errors.syntax("Syntax error"),
+          expected: "Syntax error",
+        },
+        {
+          error: MockNotebook.errors.unknown("Unknown error"),
+          expected: "Unknown error",
+        },
+      ];
+
+      for (const { error, expected } of errorTypes) {
+        // Create a fresh store and provider for each error type
+        const testStore = createStore();
+        const testProvider = new ErrorContextProvider(testStore);
+
+        const { notebookState } = MockNotebook.cellWithErrors("Cell 1", [
+          error,
+        ]);
+        testStore.set(notebookAtom, notebookState);
+
+        const items = testProvider.getItems();
+        expect(items).toHaveLength(1);
+        const context = testProvider.formatContext(items[0]);
+        expect(context).toContain(expected);
+      }
+    });
+
+    it("should handle unknown error type gracefully", () => {
+      const error = { type: "unknown-type" as any };
+      const { notebookState } = MockNotebook.cellWithErrors("Cell 1", [error]);
+      store.set(notebookAtom, notebookState);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[0]);
+      expect(context).toContain("Unknown error");
+    });
+  });
+});

package/src/core/ai/context/providers/common.ts

@@ -4,4 +4,5 @@ export const Boosts = {
   LOCAL_TABLE: 5,
   REMOTE_TABLE: 4,
   VARIABLE: 3,
+  ERROR: 2,
 } as const;

package/src/core/ai/context/providers/error.ts

@@ -0,0 +1,165 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Completion } from "@codemirror/autocomplete";
+import { cellErrorsAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import type { MarimoError } from "@/core/kernel/messages";
+import type { JotaiStore } from "@/core/state/jotai";
+import { PluralWord } from "@/utils/pluralize";
+import { type AIContextItem, AIContextProvider } from "../registry";
+import { contextToXml } from "../utils";
+import { Boosts } from "./common";
+
+export interface ErrorContextItem extends AIContextItem {
+  type: "error";
+  data: {
+    type: "all-errors";
+    errors: Array<{
+      cellId: CellId;
+      cellName: string;
+      errorData: MarimoError[];
+    }>;
+  };
+}
+
+function describeError(error: MarimoError): string {
+  if (error.type === "setup-refs") {
+    return "The setup cell cannot have references";
+  }
+  if (error.type === "cycle") {
+    return "This cell is in a cycle";
+  }
+  if (error.type === "multiple-defs") {
+    return `The variable '${error.name}' was defined by another cell`;
+  }
+  if (error.type === "import-star") {
+    return error.msg;
+  }
+  if (error.type === "ancestor-stopped") {
+    return error.msg;
+  }
+  if (error.type === "ancestor-prevented") {
+    return error.msg;
+  }
+  if (error.type === "exception") {
+    return error.msg;
+  }
+  if (error.type === "strict-exception") {
+    return error.msg;
+  }
+  if (error.type === "interruption") {
+    return "This cell was interrupted and needs to be re-run";
+  }
+  if (error.type === "syntax") {
+    return error.msg;
+  }
+  if (error.type === "unknown") {
+    return error.msg;
+  }
+  return "Unknown error";
+}
+
+const errorsTxt = new PluralWord("error", "errors");
+export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
+  readonly title = "Errors";
+  readonly mentionPrefix = "@";
+  readonly contextType = "error";
+
+  constructor(private store: JotaiStore) {
+    super();
+  }
+
+  getItems(): ErrorContextItem[] {
+    const errors = this.store.get(cellErrorsAtom);
+    if (errors.length === 0) {
+      return [];
+    }
+
+    return [
+      {
+        uri: this.asURI("all"),
+        name: "Errors",
+        type: this.contextType,
+        data: {
+          type: "all-errors",
+          errors: errors.map((error) => ({
+            cellId: error.cellId,
+            cellName: error.cellName,
+            errorData: error.output.data,
+          })),
+        },
+        description: "All errors in the notebook",
+      },
+    ];
+
+    // TODO: maybe handle single errors or grouped by types
+  }
+
+  formatCompletion(item: ErrorContextItem): Completion {
+    if (item.data.type === "all-errors") {
+      return {
+        label: "@Errors",
+        displayLabel: "Errors",
+        detail: `${item.data.errors.length} ${errorsTxt.pluralize(item.data.errors.length)}`,
+        boost: Boosts.ERROR,
+        type: "error",
+        apply: "@Errors",
+        section: "Errors",
+        info: () => {
+          const infoContainer = document.createElement("div");
+          infoContainer.classList.add(
+            "mo-cm-tooltip",
+            "docs-documentation",
+            "min-w-[200px]",
+            "flex",
+            "flex-col",
+            "gap-1",
+            "p-2",
+          );
+
+          const headerDiv = document.createElement("div");
+          headerDiv.classList.add("flex", "flex-col", "gap-1");
+
+          const nameDiv = document.createElement("div");
+          nameDiv.classList.add("font-bold", "text-base");
+          nameDiv.textContent = "Errors";
+          headerDiv.append(nameDiv);
+
+          const descriptionDiv = document.createElement("div");
+          descriptionDiv.classList.add("text-sm", "text-muted-foreground");
+          descriptionDiv.textContent = `${item.data.errors.length} ${errorsTxt.pluralize(item.data.errors.length)}`;
+          headerDiv.append(descriptionDiv);
+
+          infoContainer.append(headerDiv);
+
+          return infoContainer;
+        },
+      };
+    }
+
+    return {
+      label: "Error",
+      displayLabel: "Error",
+      boost: Boosts.ERROR,
+    };
+  }
+
+  formatContext(item: ErrorContextItem): string {
+    const { data } = item;
+
+    const xmls = data.errors.map((err) => {
+      return contextToXml({
+        type: this.contextType,
+        data: {
+          name: err.cellName || `Cell ${err.cellId}`,
+          description: err.errorData
+            .map((err) => describeError(err))
+            .join("\n"),
+        },
+      });
+    });
+
+    const xml = xmls.join("\n\n");
+    return xml;
+  }
+}

package/src/core/ai/context/providers/tables.ts

@@ -5,6 +5,7 @@ import type { DatasetTablesMap } from "@/core/datasets/data-source-connections";
 import type { DataTable } from "@/core/kernel/messages";
 import type { AIContextItem } from "../registry";
 import { AIContextProvider } from "../registry";
+import { contextToXml } from "../utils";
 import { Boosts } from "./common";
 
 export interface TableContextItem extends AIContextItem {
@@ -32,8 +33,8 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
   }
 
   formatContext(item: TableContextItem): string {
-    const { data, uri: id } = item;
-    const { columns, source, num_rows, num_columns } = data;
+    const { data } = item;
+    const { columns, source, num_rows, num_columns, name } = data;
     const shape = [
       num_rows == null ? undefined : `${num_rows} rows`,
       num_columns == null ? undefined : `${num_columns} columns`,
@@ -41,16 +42,23 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
       .filter(Boolean)
       .join(", ");
 
-    let context = `Table: ${id}\nSource: ${source || "unknown"}`;
+    let details = "";
     if (shape) {
-      context += `\nShape: ${shape}`;
+      details += `\nShape: ${shape}`;
     }
 
     if (columns && columns.length > 0) {
-      context += `\nColumns:\n${columns.map((col) => `  - ${col.name}: ${col.type}`).join("\n")}`;
+      details += `\nColumns:\n${columns.map((col) => `  - ${col.name}: ${col.type}`).join("\n")}`;
     }
 
-    return context;
+    return contextToXml({
+      type: this.contextType,
+      data: {
+        name: name,
+        source: source ?? "unknown",
+      },
+      details: details,
+    });
   }
 
   formatCompletion(item: TableContextItem): Completion {

package/src/core/ai/context/providers/variable.ts

@@ -5,6 +5,7 @@ import { createVariableInfoElement } from "@/core/codemirror/completion/variable
 import type { DatasetTablesMap } from "@/core/datasets/data-source-connections";
 import type { Variable, Variables } from "@/core/variables/types";
 import { type AIContextItem, AIContextProvider } from "../registry";
+import { contextToXml } from "../utils";
 import { Boosts } from "./common";
 
 export interface VariableContextItem extends AIContextItem {
@@ -65,8 +66,16 @@ export class VariableContextProvider extends AIContextProvider<VariableContextIt
   }
 
   formatContext(item: VariableContextItem): string {
-    const { uri: id, data } = item;
+    const { data } = item;
     const { variable } = data;
-    return `Variable: ${id}\nType: ${variable.dataType || "unknown"}\nPreview: ${JSON.stringify(variable.value)}`;
+    return contextToXml({
+      type: this.contextType,
+      data: {
+        name: variable.name,
+        dataType: variable.dataType,
+      },
+      details:
+        variable.value == null ? undefined : JSON.stringify(variable.value),
+    });
   }
 }

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

@@ -0,0 +1,49 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+export interface AiContextPayload {
+  type: string;
+  data: Record<string, unknown>;
+  details?: string;
+}
+
+// XML escaping utility
+function escapeXml(unsafe: string): string {
+  return (
+    unsafe
+      // We don't escape these characters because this is for an LLM and they can interpret this just fine.
+      // .replaceAll("&", "&amp;")
+      // .replaceAll('"', "&quot;")
+      // .replaceAll("'", "&#39;")
+      .replaceAll("<", "&lt;")
+      .replaceAll(">", "&gt;")
+  );
+}
+
+// Convert a single Context object to XML
+export function contextToXml(context: AiContextPayload): string {
+  const { type, data, details } = context;
+
+  // Start with opening tag
+  let xml = `<${type}`;
+
+  // Add data as attributes
+  for (const [key, value] of Object.entries(data)) {
+    const escapedValue = escapeXml(String(value));
+    if (value !== undefined) {
+      xml += ` ${key}="${escapedValue}"`;
+    }
+  }
+
+  // Close the opening tag
+  xml += ">";
+
+  // Add details as content if present
+  if (details) {
+    xml += escapeXml(details);
+  }
+
+  // Add closing tag
+  xml += `</${type}>`;
+
+  return xml;
+}

package/src/core/cells/cells.ts

@@ -1489,7 +1489,7 @@ export const canUndoDeletesAtom = atom((get) =>
 
 export const needsRunAtom = atom((get) => notebookNeedsRun(get(notebookAtom)));
 
-const cellErrorsAtom = atom((get) => {
+export const cellErrorsAtom = atom((get) => {
   const { cellIds, cellRuntime, cellData } = get(notebookAtom);
   const errors = cellIds.inOrderIds
     .map((cellId) => {

package/dist/assets/_baseMap-u5TY8JMW.js

@@ -1 +0,0 @@
-import{b as m}from"./_baseEach-DJPM6ZzY.js";import{x as s}from"./index-CQFSXxPp.js";function e(r,o){var a=-1,t=s(r)?Array(r.length):[];return m(r,function(n,f,i){t[++a]=o(n,f,i)}),t}export{e as b};

package/dist/assets/channel-BIZI32eZ.js

@@ -1 +0,0 @@
-import{U as s,C as o}from"./mermaid-CUkYs7pX.js";const n=(a,r)=>s.lang.round(o.parse(a)[r]);export{n as c};

package/dist/assets/classDiagram-3BZAVTQC-DfPpF-gd.js

@@ -1 +0,0 @@
-import{s as a,c as s,a as t,C as o}from"./chunk-JBRWN2VN-CB8a_SQv.js";import{_ as e}from"./mermaid-CUkYs7pX.js";import"./transform-B8bpuzxV.js";import"./chunk-GLLZNHP4-CRWkmhQj.js";import"./chunk-WVR4S24B-ZnYlHppS.js";import"./chunk-NRVI72HA-nFtyi9pH.js";import"./index-CQFSXxPp.js";import"./step-BwsUM5iJ.js";import"./timer-BwIYMJWC.js";var i={parser:t,get db(){return new o},renderer:s,styles:a,init:e(r=>{r.class||(r.class={}),r.class.arrowMarkerAbsolute=r.arrowMarkerAbsolute},"init")};export{i as diagram};

package/dist/assets/classDiagram-v2-QTMF73CY-DfPpF-gd.js

@@ -1 +0,0 @@
-import{s as a,c as s,a as t,C as o}from"./chunk-JBRWN2VN-CB8a_SQv.js";import{_ as e}from"./mermaid-CUkYs7pX.js";import"./transform-B8bpuzxV.js";import"./chunk-GLLZNHP4-CRWkmhQj.js";import"./chunk-WVR4S24B-ZnYlHppS.js";import"./chunk-NRVI72HA-nFtyi9pH.js";import"./index-CQFSXxPp.js";import"./step-BwsUM5iJ.js";import"./timer-BwIYMJWC.js";var i={parser:t,get db(){return new o},renderer:s,styles:a,init:e(r=>{r.class||(r.class={}),r.class.arrowMarkerAbsolute=r.arrowMarkerAbsolute},"init")};export{i as diagram};

package/dist/assets/clone-yHYNh3QG.js

@@ -1 +0,0 @@
-import{b as n}from"./_baseUniq-CrXjB6JP.js";function o(r){return n(r,4)}export{o as c};

package/dist/assets/stateDiagram-v2-EYPG3UTE-DX5NRIHK.js

@@ -1 +0,0 @@
-import{s as t,b as a,a as e,S as o}from"./chunk-LXBSTHXV-YNRk01ob.js";import{_ as s}from"./mermaid-CUkYs7pX.js";import"./transform-B8bpuzxV.js";import"./chunk-WVR4S24B-ZnYlHppS.js";import"./chunk-NRVI72HA-nFtyi9pH.js";import"./index-CQFSXxPp.js";import"./step-BwsUM5iJ.js";import"./timer-BwIYMJWC.js";var i={parser:e,get db(){return new o(2)},renderer:a,styles:t,init:s(r=>{r.state||(r.state={}),r.state.arrowMarkerAbsolute=r.arrowMarkerAbsolute},"init")};export{i as diagram};