@marimo-team/frontend 0.15.4 → 0.16.0-dev96986

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

@@ -8,7 +8,7 @@ import type { DatasetsState } from "@/core/datasets/types";
 import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import type { Variable, VariableName } from "@/core/variables/types";
-import { getAICompletionBody } from "../completion-utils";
+import { codeToCells, getAICompletionBody } from "../completion-utils";
 
 // Mock getCodes function
 vi.mock("@/core/codemirror/copilot/getCodes", () => ({
@@ -357,3 +357,180 @@ describe("getAICompletionBody", () => {
     `);
   });
 });
+
+describe("codeToCells", () => {
+  it("should return empty array for empty string", () => {
+    const code = "";
+    const result = codeToCells(code);
+    expect(result).toEqual([]);
+  });
+
+  it("should return empty array for whitespace only", () => {
+    const code = "   \n\t  ";
+    const result = codeToCells(code);
+    expect(result).toEqual([]);
+  });
+
+  it("should convert code without backticks to single python cell", () => {
+    const code = "print('Hello, world!')";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should convert code with single closed backticks to cells", () => {
+    const code = "```python\nprint('Hello, world!')\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should convert code with unclosed backticks to cells", () => {
+    const code = "```python\nprint('Hello, world!')\n";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should convert code with multiple closed cells", () => {
+    const code =
+      "```python\nprint('Hello, world!')\n```\n```sql\nSELECT * FROM users\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+      { language: "sql", code: "SELECT * FROM users" },
+    ]);
+  });
+
+  it("should handle code with no language identifier", () => {
+    const code = "```\nprint('Hello, world!')\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should handle unclosed code with no language identifier", () => {
+    const code = "```\nprint('Hello, world!')\n";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should handle markdown language", () => {
+    const code = "```markdown\n# Hello, world!\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([{ language: "markdown", code: "# Hello, world!" }]);
+  });
+
+  it("should handle sql language", () => {
+    const code = "```sql\nSELECT * FROM users\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([{ language: "sql", code: "SELECT * FROM users" }]);
+  });
+
+  it("should handle unclosed markdown cell", () => {
+    const code = "```markdown\n# Hello, world!\n";
+    const result = codeToCells(code);
+    expect(result).toEqual([{ language: "markdown", code: "# Hello, world!" }]);
+  });
+
+  it("should handle unclosed sql cell", () => {
+    const code = "```sql\nSELECT * FROM users\n";
+    const result = codeToCells(code);
+    expect(result).toEqual([{ language: "sql", code: "SELECT * FROM users" }]);
+  });
+
+  it("should handle empty cells and skip them", () => {
+    const code = "```python\n\n```\n```sql\nSELECT * FROM users\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([{ language: "sql", code: "SELECT * FROM users" }]);
+  });
+
+  it("should handle cells with only whitespace and skip them", () => {
+    const code = "```python\n   \n```\n```sql\nSELECT * FROM users\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([{ language: "sql", code: "SELECT * FROM users" }]);
+  });
+
+  it("should handle code with trailing newlines", () => {
+    const code = "```python\nprint('Hello, world!')\n\n\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should handle unclosed code with trailing newlines", () => {
+    const code = "```python\nprint('Hello, world!')\n\n\n";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should handle multiple cells with different languages", () => {
+    const code =
+      "```python\nprint('Hello, world!')\n```\n```sql\nSELECT * FROM users\n```\n```markdown\n# Title\nThis is markdown\n```";
+
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+      { language: "sql", code: "SELECT * FROM users" },
+      { language: "markdown", code: "# Title\nThis is markdown" },
+    ]);
+  });
+
+  it("should handle complex multiline code", () => {
+    const code =
+      '```python\ndef hello():\n    print("Hello, world!")\n    return "success"\n\nhello()\n```\n```sql\nSELECT \n    id,\n    name,\n    email\nFROM users\nWHERE active = true\nORDER BY name;\n```';
+
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      {
+        language: "python",
+        code: 'def hello():\n    print("Hello, world!")\n    return "success"\n\nhello()',
+      },
+      {
+        language: "sql",
+        code: "SELECT \n    id,\n    name,\n    email\nFROM users\nWHERE active = true\nORDER BY name;",
+      },
+    ]);
+  });
+
+  it("should handle code with backticks in the content", () => {
+    const code = "```python\nprint('```')\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([{ language: "python", code: "print('" }]);
+  });
+
+  it("should handle code with no backticks in the last cell", () => {
+    const code =
+      "```python\nprint('Hello, world!')\n```\n```python\nprint('Hello, world!')";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should handle case insensitive language detection", () => {
+    const code = "```PYTHON\nprint('Hello, world!')\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "python", code: "print('Hello, world!')" },
+    ]);
+  });
+
+  it("should handle unknown language", { fails: true }, () => {
+    const code = "```javascript\nconsole.log('Hello, world!')\n```";
+    const result = codeToCells(code);
+    expect(result).toEqual([
+      { language: "javascript", code: "console.log('Hello, world!')" },
+    ]);
+  });
+});

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -1,6 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import { useCompletion } from "@ai-sdk/react";
+import { useChat } from "@ai-sdk/react";
 import {
   autocompletion,
   type Completion,
@@ -8,7 +8,6 @@ import {
   type CompletionSource,
 } from "@codemirror/autocomplete";
 import { markdown } from "@codemirror/lang-markdown";
-import { sql } from "@codemirror/lang-sql";
 import { Prec } from "@codemirror/state";
 import { promptHistory, storePrompt } from "@marimo-team/codemirror-ai";
 import ReactCodeMirror, {
@@ -17,7 +16,7 @@ import ReactCodeMirror, {
   minimalSetup,
   type ReactCodeMirrorRef,
 } from "@uiw/react-codemirror";
-import { useAtom, useStore } from "jotai";
+import { useAtom, useAtomValue, useStore } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import {
   ChevronsUpDown,
@@ -31,6 +30,10 @@ import { useMemo, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
 import { z } from "zod";
 import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
+import {
+  buildCompletionRequestBody,
+  handleToolCall,
+} from "@/components/chat/chat-utils";
 import { Button } from "@/components/ui/button";
 import {
   DropdownMenu,
@@ -40,31 +43,22 @@ import {
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
 import { toast } from "@/components/ui/use-toast";
+import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
 import { resourceExtension } from "@/core/codemirror/ai/resources";
-import { customPythonLanguageSupport } from "@/core/codemirror/language/languages/python";
-import { SQLLanguageAdapter } from "@/core/codemirror/language/languages/sql/sql";
+import { useRequestClient } from "@/core/network/requests";
+import type { AiCompletionRequest } from "@/core/network/types";
 import { useRuntimeManager } from "@/core/runtime/config";
 import { useTheme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { prettyError } from "@/utils/errors";
 import { ZodLocalStorage } from "@/utils/localStorage";
-import { useCellActions } from "../../../core/cells/cells";
 import { PythonIcon } from "../cell/code/icons";
 import {
   CompletionActions,
   createAiCompletionOnKeydown,
 } from "./completion-handlers";
-import {
-  CONTEXT_TRIGGER,
-  getAICompletionBody,
-  mentionsCompletionSource,
-} from "./completion-utils";
-
-const pythonExtensions = [
-  customPythonLanguageSupport(),
-  EditorView.lineWrapping,
-];
-const sqlExtensions = [sql(), EditorView.lineWrapping];
+import { CONTEXT_TRIGGER, mentionsCompletionSource } from "./completion-utils";
+import { StreamingChunkTransport } from "./transport/chat-transport";
 
 // Persist across sessions
 const languageAtom = atomWithStorage<"python" | "sql">(
@@ -82,32 +76,56 @@ const promptHistoryStorage = new ZodLocalStorage(z.array(z.string()), () => []);
 export const AddCellWithAI: React.FC<{
   onClose: () => void;
 }> = ({ onClose }) => {
-  const { createNewCell } = useCellActions();
-  const [completionBody, setCompletionBody] = useState<object>({});
+  const store = useStore();
+  const [input, setInput] = useState("");
+
+  const { deleteAllStagedCells, clearStagedCells, onStream } =
+    useStagedCells(store);
   const [language, setLanguage] = useAtom(languageAtom);
-  const { theme } = useTheme();
   const runtimeManager = useRuntimeManager();
+  const { invokeAiTool } = useRequestClient();
 
+  const stagedAICells = useAtomValue(stagedAICellsAtom);
   const inputRef = useRef<ReactCodeMirrorRef>(null);
 
-  const {
-    completion,
-    input,
-    stop,
-    isLoading,
-    setCompletion,
-    setInput,
-    handleSubmit,
-  } = useCompletion({
-    api: runtimeManager.getAiURL("completion").toString(),
-    headers: runtimeManager.headers(),
-    streamProtocol: "text",
+  const { sendMessage, stop, status, addToolResult } = useChat({
     // Throttle the messages and data updates to 100ms
     experimental_throttle: 100,
-    body: {
-      ...completionBody,
-      language: language,
-      code: "",
+    transport: new StreamingChunkTransport(
+      {
+        api: runtimeManager.getAiURL("completion").toString(),
+        headers: runtimeManager.headers(),
+        prepareSendMessagesRequest: async (options) => {
+          const completionBody = await buildCompletionRequestBody(
+            options.messages,
+          );
+          const body: AiCompletionRequest = {
+            ...options,
+            ...completionBody,
+            code: "",
+            prompt: "", // Don't need prompt since we are using messages
+            language: language,
+          };
+
+          return {
+            body: body,
+          };
+        },
+      },
+      (chunk) => {
+        onStream(chunk);
+      },
+    ),
+    onToolCall: async ({ toolCall }) => {
+      await handleToolCall({
+        invokeAiTool,
+        addToolResult,
+        toolCall: {
+          toolName: toolCall.toolName,
+          toolCallId: toolCall.toolCallId,
+          input: toolCall.input as Record<string, never>,
+        },
+      });
     },
     onError: (error) => {
       toast({
@@ -115,18 +133,20 @@ export const AddCellWithAI: React.FC<{
         description: prettyError(error),
       });
     },
-    onFinish: (_prompt, completion) => {
-      // Remove trailing new lines
-      setCompletion(completion.trimEnd());
-    },
   });
 
+  const isLoading = status === "streaming" || status === "submitted";
+  const hasCompletion = stagedAICells.size > 0;
+  const multipleCompletions = stagedAICells.size > 1;
+
   const submit = () => {
     if (!isLoading) {
       if (inputRef.current?.view) {
         storePrompt(inputRef.current.view);
       }
-      handleSubmit();
+      // TODO: When we have conversations, don't delete existing cells
+      deleteAllStagedCells();
+      sendMessage({ text: input });
     }
   };
 
@@ -171,20 +191,12 @@ export const AddCellWithAI: React.FC<{
   );
 
   const handleAcceptCompletion = () => {
-    createNewCell({
-      cellId: "__end__",
-      before: false,
-      code:
-        language === "python"
-          ? completion
-          : SQLLanguageAdapter.fromQuery(completion),
-    });
-    setCompletion("");
+    clearStagedCells();
     onClose();
   };
 
   const handleDeclineCompletion = () => {
-    setCompletion("");
+    deleteAllStagedCells();
   };
 
   const inputComponent = (
@@ -193,20 +205,19 @@ export const AddCellWithAI: React.FC<{
       <PromptInput
         inputRef={inputRef}
         onClose={() => {
-          setCompletion("");
+          deleteAllStagedCells();
           onClose();
         }}
         value={input}
         onChange={(newValue) => {
           setInput(newValue);
-          setCompletionBody(getAICompletionBody({ input: newValue }));
         }}
         onSubmit={submit}
         onKeyDown={createAiCompletionOnKeydown({
           handleAcceptCompletion,
           handleDeclineCompletion,
           isLoading,
-          completion,
+          hasCompletion,
         })}
       />
       {isLoading && (
@@ -234,7 +245,7 @@ export const AddCellWithAI: React.FC<{
     <div className={cn("flex flex-col w-full gap-2 py-2")}>
       {inputComponent}
       <div className="flex flex-row justify-between -mt-1 ml-1 mr-3">
-        {!completion && (
+        {!hasCompletion && (
           <span className="text-xs text-muted-foreground px-3 flex flex-col gap-1">
             <span>
               You can mention{" "}
@@ -245,12 +256,13 @@ export const AddCellWithAI: React.FC<{
             <span>Code from other cells is automatically included.</span>
           </span>
         )}
-        {completion && (
+        {hasCompletion && (
           <CompletionActions
             isLoading={isLoading}
             onAccept={handleAcceptCompletion}
             onDecline={handleDeclineCompletion}
             size="sm"
+            multipleCompletions={multipleCompletions}
           />
         )}
         <div className="ml-auto flex items-center gap-1">
@@ -262,16 +274,6 @@ export const AddCellWithAI: React.FC<{
           />
         </div>
       </div>
-
-      {completion && (
-        <ReactCodeMirror
-          value={completion}
-          className="cm border-t"
-          onChange={setCompletion}
-          theme={theme === "dark" ? "dark" : "light"}
-          extensions={language === "python" ? pythonExtensions : sqlExtensions}
-        />
-      )}
     </div>
   );
 };

package/src/components/editor/ai/ai-completion-editor.tsx

@@ -165,7 +165,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
             <SparklesIcon className="text-(--blue-10) shrink-0" size={16} />
             <PromptInput
               inputRef={inputRef}
-              className="h-full my-0 py-2"
+              className="h-full my-0 py-2 flex items-center"
               onClose={() => {
                 declineChange();
                 setCompletion("");
@@ -187,7 +187,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
                 handleAcceptCompletion,
                 handleDeclineCompletion,
                 isLoading,
-                completion,
+                hasCompletion: completion.trim().length > 0,
               })}
             />
             {isLoading && (
@@ -202,30 +202,33 @@ export const AiCompletionEditor: React.FC<Props> = ({
                 Stop
               </Button>
             )}
-            {completion && (
-              <CompletionActions
-                isLoading={isLoading}
-                onAccept={handleAcceptCompletion}
-                onDecline={handleDeclineCompletion}
-                size="xs"
-              />
-            )}
-            <div className="flex flex-row items-center gap-0.5 -ml-1.5 -mr-2">
-              <Tooltip content="Add context">
-                <Button
-                  variant="text"
-                  size="icon"
-                  onClick={() => addContextCompletion(inputRef)}
-                >
-                  <AtSignIcon className="h-3 w-3" />
-                </Button>
-              </Tooltip>
-              <AIModelDropdown
-                triggerClassName="h-7"
-                iconSize="small"
-                displayIconOnly={true}
-                forRole="edit"
-              />
+            <div className="-mr-1.5 py-1.5">
+              <div className="flex flex-row items-center justify-end gap-0.5">
+                <Tooltip content="Add context">
+                  <Button
+                    variant="text"
+                    size="icon"
+                    onClick={() => addContextCompletion(inputRef)}
+                  >
+                    <AtSignIcon className="h-3 w-3" />
+                  </Button>
+                </Tooltip>
+                <AIModelDropdown
+                  triggerClassName="h-7 text-xs w-24"
+                  iconSize="small"
+                  forRole="edit"
+                />
+              </div>
+              {completion && (
+                <div className="-mb-1.5">
+                  <CompletionActions
+                    isLoading={isLoading}
+                    onAccept={handleAcceptCompletion}
+                    onDecline={handleDeclineCompletion}
+                    size="xs"
+                  />
+                </div>
+              )}
             </div>
 
             <div className="h-full w-px bg-border mx-2" />

package/src/components/editor/ai/completion-handlers.tsx

@@ -12,20 +12,20 @@ export const createAiCompletionOnKeydown = (opts: {
   handleAcceptCompletion: () => void;
   handleDeclineCompletion: () => void;
   isLoading: boolean;
-  completion: string;
+  hasCompletion: boolean;
 }) => {
   const {
     handleAcceptCompletion,
     handleDeclineCompletion,
     isLoading,
-    completion,
+    hasCompletion,
   } = opts;
 
   return (e: React.KeyboardEvent<HTMLDivElement>) => {
     const metaKey = isPlatformMac() ? e.metaKey : e.ctrlKey;
 
     // Mod+Enter should accept the completion, if there is one
-    if (metaKey && e.key === "Enter" && !isLoading && completion) {
+    if (metaKey && e.key === "Enter" && !isLoading && hasCompletion) {
       handleAcceptCompletion();
     }
 
@@ -49,6 +49,7 @@ export const CompletionActions: React.FC<{
   size?: "xs" | "sm";
   acceptShortcut?: string;
   declineShortcut?: string;
+  multipleCompletions?: boolean;
 }> = ({
   isLoading,
   onAccept,
@@ -56,6 +57,7 @@ export const CompletionActions: React.FC<{
   size = "sm",
   acceptShortcut = "Mod-↵",
   declineShortcut = "Shift-Mod-Delete",
+  multipleCompletions = false,
 }) => {
   return (
     <>
@@ -68,7 +70,7 @@ export const CompletionActions: React.FC<{
         onClick={onAccept}
       >
         <span className="text-(--grass-11) opacity-100">
-          Accept{" "}
+          Accept{multipleCompletions && " all"}{" "}
           <MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
         </span>
       </Button>
@@ -76,14 +78,50 @@ export const CompletionActions: React.FC<{
         data-testid="decline-completion-button"
         variant="text"
         size={size}
-        className="mb-0 pl-1"
+        className="mb-0 pl-1 pr-0"
         onClick={onDecline}
       >
         <span className="text-(--red-10)">
-          Reject{" "}
+          Reject{multipleCompletions && " all"}{" "}
           <MinimalHotkeys className="ml-1 inline" shortcut={declineShortcut} />
         </span>
       </Button>
     </>
   );
 };
+
+export const CompletionActionsCellFooter: React.FC<{
+  isLoading: boolean;
+  onAccept: () => void;
+  onDecline: () => void;
+  size?: "xs" | "sm";
+  multipleCompletions?: boolean;
+}> = ({ isLoading, onAccept, onDecline }) => {
+  return (
+    <>
+      <Button
+        variant="text"
+        size="xs"
+        disabled={isLoading}
+        onClick={onAccept}
+        className="h-6 text-(--grass-11) border-(--grass-7) bg-(--grass-3)/60
+        hover:bg-(--grass-3) dark:bg-(--grass-4)/80 dark:hover:bg-(--grass-3) rounded px-3 font-semibold
+        active:bg-(--grass-5) dark:active:bg-(--grass-4)
+        border-(--green-6) border hover:shadow-xs"
+      >
+        Accept
+      </Button>
+      <Button
+        variant="text"
+        size="xs"
+        onClick={onDecline}
+        className="h-6 text-(--red-10) border-(--red-7) bg-(--red-3)/60 hover:bg-(--red-3)
+        dark:bg-(--red-4)/80 dark:hover:bg-(--red-3) rounded px-3 font-semibold
+        active:bg-(--red-5) dark:active:bg-(--red-4)
+        border-(--red-6) border hover:shadow-xs"
+      >
+        Reject
+      </Button>
+    </>
+  );
+};