@marimo-team/frontend 0.15.1-dev33 → 0.15.1-dev35

package/src/components/editor/ai/__tests__/completion-utils.test.ts

@@ -52,41 +52,40 @@ describe("getAICompletionBody", () => {
     const input = "Use @data://dataset1 and @data://dataset2 for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "dataset1",
-            columns: [
-              { name: "col1", type: "number" },
-              { name: "col2", type: "string" },
-            ],
-          },
-          {
-            name: "dataset2",
-            columns: [
-              { name: "col3", type: "boolean" },
-              { name: "col4", type: "date" },
-            ],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="dataset1" source="unknown">
+      Columns:
+        - col1: number
+        - col2: string</data>
+
+      <data name="dataset2" source="unknown">
+      Columns:
+        - col3: boolean
+        - col4: date</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle input with no mentioned datasets", () => {
     const input = "Perform some analysis without mentioning @data://datasets";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle input with non-existent datasets", () => {
@@ -106,21 +105,19 @@ describe("getAICompletionBody", () => {
       "Use @data://existingDataset and @data://nonExistentDataset for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "existingDataset",
-            columns: [
-              { name: "col1", type: "number" },
-              { name: "col2", type: "string" },
-            ],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="existingDataset" source="unknown">
+      Columns:
+        - col1: number
+        - col2: string</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle dataset names with dots", () => {
@@ -144,25 +141,23 @@ describe("getAICompletionBody", () => {
       "Use @data://dataset.with.dots and @data://regular_dataset for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "dataset.with.dots",
-            columns: [
-              { name: "col1", type: "number" },
-              { name: "col2", type: "string" },
-            ],
-          },
-          {
-            name: "regular_dataset",
-            columns: [{ name: "col3", type: "boolean" }],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="dataset.with.dots" source="unknown">
+      Columns:
+        - col1: number
+        - col2: string</data>
+
+      <data name="regular_dataset" source="unknown">
+      Columns:
+        - col3: boolean</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle connections", () => {
@@ -200,18 +195,18 @@ describe("getAICompletionBody", () => {
     const input = "Use @data://table1 for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "table1",
-            columns: [{ name: "col1", type: "number" }],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="table1" source="unknown">
+      Columns:
+        - col1: number</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should return the correct completion body with mentioned variables", () => {
@@ -237,24 +232,18 @@ describe("getAICompletionBody", () => {
     const input = "Use @variable://var1 and @variable://var2 for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [],
-        variables: [
-          {
-            name: "var1",
-            valueType: "string",
-            previewValue: "string value",
-          },
-          {
-            name: "var2",
-            valueType: "number",
-            previewValue: "42",
-          },
-        ],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<variable name="var1" dataType="string">"string value"</variable>
+
+      <variable name="var2" dataType="number">"42"</variable>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle input with both datasets and variables", () => {
@@ -284,27 +273,21 @@ describe("getAICompletionBody", () => {
     const input = "Use @data://dataset1 and @variable://var1 for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "dataset1",
-            columns: [
-              { name: "col1", type: "number" },
-              { name: "col2", type: "string" },
-            ],
-          },
-        ],
-        variables: [
-          {
-            name: "var1",
-            valueType: "string",
-            previewValue: "string value",
-          },
-        ],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="dataset1" source="unknown">
+      Columns:
+        - col1: number
+        - col2: string</data>
+
+      <variable name="var1" dataType="string">"string value"</variable>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle non-existent variables", () => {
@@ -324,19 +307,16 @@ describe("getAICompletionBody", () => {
       "Use @variable://existingVar and @variable://nonExistentVar for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [],
-        variables: [
-          {
-            name: "existingVar",
-            valueType: "string",
-            previewValue: "string value",
-          },
-        ],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<variable name="existingVar" dataType="string">"string value"</variable>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should prioritize datasets over variables when there's a name conflict", () => {
@@ -363,17 +343,17 @@ describe("getAICompletionBody", () => {
     const input = "Use @data://conflict for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "conflict",
-            columns: [{ name: "col1", type: "number" }],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="conflict" source="unknown">
+      Columns:
+        - col1: number</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 });

package/src/components/editor/ai/completion-utils.ts

@@ -6,13 +6,9 @@ import type {
   CompletionSource,
 } from "@codemirror/autocomplete";
 import { getAIContextRegistry } from "@/core/ai/context/context";
-import type { TableContextItem } from "@/core/ai/context/providers/tables";
-import type { VariableContextItem } from "@/core/ai/context/providers/variable";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
-import type { DataTable } from "@/core/kernel/messages";
 import type { AiCompletionRequest } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
-import type { Variable } from "@/core/variables/types";
 import { Logger } from "@/utils/Logger";
 
 /**
@@ -23,26 +19,15 @@ export function getAICompletionBody({
 }: {
   input: string;
 }): Omit<AiCompletionRequest, "language" | "prompt" | "code"> {
-  const { datasets, variables } = extractDatasetsAndVariables(input);
-  Logger.debug("Included datasets", datasets);
-  Logger.debug("Included variables", variables);
+  const contextString = extractDatasetsAndVariables(input);
+  Logger.debug("Included context", contextString);
 
   return {
     includeOtherCode: getCodes(""),
     context: {
-      schema: datasets.map((dataset) => ({
-        name: dataset.name,
-        columns: dataset.columns.map((column) => ({
-          name: column.name,
-          type: column.type,
-          sampleValues: column.sample_values,
-        })),
-      })),
-      variables: variables.map((variable) => ({
-        name: variable.name,
-        valueType: variable.dataType ?? "",
-        previewValue: variable.value,
-      })),
+      plainText: contextString,
+      schema: [],
+      variables: [],
     },
   };
 }
@@ -52,21 +37,10 @@ export function getAICompletionBody({
  * References are with @<name> in the input.
  * Prioritizes datasets over variables if there's a name conflict.
  */
-function extractDatasetsAndVariables(input: string): {
-  datasets: DataTable[];
-  variables: Variable[];
-} {
+function extractDatasetsAndVariables(input: string): string {
   const registry = getAIContextRegistry(store);
   const contextIds = registry.parseAllContextIds(input);
-  const contextInfo = registry.getContextInfo(contextIds);
-  const datasets: DataTable[] = contextInfo
-    .filter((info): info is TableContextItem => info.type === "data")
-    .map((info) => info.data);
-  const variables: Variable[] = contextInfo
-    .filter((info): info is VariableContextItem => info.type === "variable")
-    .map((info) => info.data.variable);
-
-  return { datasets, variables };
+  return registry.formatContextForAI(contextIds);
 }
 
 /**

package/src/core/ai/context/__tests__/registry.test.ts

@@ -1,7 +1,15 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import type { Completion } from "@codemirror/autocomplete";
+import { createStore } from "jotai";
 import { beforeEach, describe, expect, it } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { notebookAtom } from "@/core/cells/cells";
+import { CellId as CellIdClass } from "@/core/cells/ids";
+import {
+  type ErrorContextItem,
+  ErrorContextProvider,
+} from "../providers/error";
 import {
   type AIContextItem,
   AIContextProvider,
@@ -513,4 +521,150 @@ describe("AIContextRegistry", () => {
       expect(provider).toBe(mockProvider);
     });
   });
+
+  describe("Integration with ErrorContextProvider", () => {
+    let errorProvider: ErrorContextProvider;
+    let store: ReturnType<typeof createStore>;
+    let registry: AIContextRegistry<ErrorContextItem>;
+
+    beforeEach(() => {
+      store = createStore();
+
+      // Create mock notebook with errors using the new MockNotebook utilities
+      const cellId1 = CellIdClass.create();
+      const cellId2 = CellIdClass.create();
+
+      const notebookState = MockNotebook.notebookStateWithErrors([
+        {
+          cellId: cellId1,
+          cellName: "Cell 1",
+          errorData: [
+            MockNotebook.errors.syntax("Invalid syntax"),
+            MockNotebook.errors.exception("Runtime error"),
+          ],
+        },
+        {
+          cellId: cellId2,
+          cellName: "Cell 2",
+          errorData: [MockNotebook.errors.cycle()],
+        },
+      ]);
+
+      store.set(notebookAtom, notebookState);
+      errorProvider = new ErrorContextProvider(store);
+      registry = new AIContextRegistry<ErrorContextItem>().register(
+        errorProvider,
+      );
+    });
+
+    it("should register and provide error context items", () => {
+      const provider = registry.getProvider("error");
+      expect(provider).toBe(errorProvider);
+
+      const items = registry.getAllItems();
+      expect(items).toHaveLength(1);
+      expect(items[0].type).toBe("error");
+      expect(items[0].name).toBe("Errors");
+    });
+
+    it("should parse error context IDs correctly", () => {
+      const contextIds = registry.parseAllContextIds(
+        "Use @error://all to analyze errors",
+      );
+      expect(contextIds).toEqual(["error://all"]);
+
+      const contextInfo = registry.getContextInfo(contextIds);
+      expect(contextInfo).toHaveLength(1);
+      expect(contextInfo[0].type).toBe("error");
+    });
+
+    it("should format error context for AI", () => {
+      const contextIds = ["error://all"] as ContextLocatorId[];
+      const formattedContext = registry.formatContextForAI(contextIds);
+
+      expect(formattedContext).toContain("<error");
+      expect(formattedContext).toContain("Cell 1");
+      expect(formattedContext).toContain("Cell 2");
+      expect(formattedContext).toContain("Invalid syntax");
+      expect(formattedContext).toContain("This cell is in a cycle");
+    });
+
+    it("should provide error completions", () => {
+      // Get the error provider and test its completion directly
+      const provider = registry.getProvider("error");
+      expect(provider).toBeDefined();
+
+      const items = provider!.getItems();
+      expect(items).toHaveLength(1);
+
+      const completion = provider!.formatCompletion(items[0]);
+      expect(completion.label).toBe("@Errors");
+      expect(completion.type).toBe("error");
+    });
+
+    it("should handle empty errors gracefully", () => {
+      // Create store with no errors using MockNotebook
+      const emptyStore = createStore();
+      const emptyNotebookState = MockNotebook.notebookStateWithErrors([]);
+      emptyStore.set(notebookAtom, emptyNotebookState);
+
+      const emptyErrorProvider = new ErrorContextProvider(emptyStore);
+      const emptyRegistry = new AIContextRegistry().register(
+        emptyErrorProvider,
+      );
+
+      const items = emptyRegistry.getAllItems();
+      expect(items).toHaveLength(0);
+    });
+
+    it("should work with mixed providers", () => {
+      const mixedRegistry = new AIContextRegistry<
+        ErrorContextItem | MockContextItem
+      >()
+        .register(errorProvider)
+        .register(mockProvider);
+
+      const items = mixedRegistry.getAllItems();
+      expect(items.length).toBeGreaterThan(1);
+
+      // Should have both error and mock items
+      const errorItems = items.filter((item) => item.type === "error");
+      const mockItems = items.filter((item) => item.type === "mock");
+
+      expect(errorItems).toHaveLength(1);
+      expect(mockItems).toHaveLength(3);
+    });
+
+    it("should parse mixed context IDs", () => {
+      const mixedRegistry = new AIContextRegistry<
+        ErrorContextItem | MockContextItem
+      >()
+        .register(errorProvider)
+        .register(mockProvider);
+
+      const input = "Use @error://all and @mock://item1 for analysis";
+      const contextIds = mixedRegistry.parseAllContextIds(input);
+
+      expect(contextIds).toContain("error://all");
+      expect(contextIds).toContain("mock://item1");
+      expect(contextIds).toHaveLength(2);
+    });
+
+    it("should format mixed context for AI", () => {
+      const mixedRegistry = new AIContextRegistry<
+        ErrorContextItem | MockContextItem
+      >()
+        .register(errorProvider)
+        .register(mockProvider);
+
+      const contextIds = ["error://all", "mock://item1"] as ContextLocatorId[];
+      const formattedContext = mixedRegistry.formatContextForAI(contextIds);
+
+      // Should contain both error and mock context
+      expect(formattedContext).toContain("<error");
+      expect(formattedContext).toContain("Mock:");
+      expect(formattedContext).toContain("Cell 1");
+      expect(formattedContext).toContain("Item 1");
+    });
+  });
 });