@marimo-team/frontend 0.14.18-dev4 → 0.14.18-dev40

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

@@ -0,0 +1,423 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import {
+  autocompletion,
+  type Completion,
+  type CompletionContext,
+  type CompletionSource,
+} from "@codemirror/autocomplete";
+import { markdown } from "@codemirror/lang-markdown";
+import { sql } from "@codemirror/lang-sql";
+import { Prec } from "@codemirror/state";
+import ReactCodeMirror, {
+  EditorView,
+  keymap,
+  minimalSetup,
+  type ReactCodeMirrorRef,
+} from "@uiw/react-codemirror";
+import { useCompletion } from "ai/react";
+import { useAtom, useAtomValue, useStore } from "jotai";
+import { atomWithStorage } from "jotai/utils";
+import {
+  ChevronsUpDown,
+  DatabaseIcon,
+  Loader2Icon,
+  SendHorizontal,
+  SparklesIcon,
+  XIcon,
+} from "lucide-react";
+import { useMemo, useState } from "react";
+import useEvent from "react-use-event-hook";
+import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
+import { Button } from "@/components/ui/button";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuSeparator,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { toast } from "@/components/ui/use-toast";
+import { useModelChange } from "@/core/ai/config";
+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 { aiAtom } from "@/core/config/config";
+import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
+import { useRuntimeManager } from "@/core/runtime/config";
+import { useTheme } from "@/theme/useTheme";
+import { cn } from "@/utils/cn";
+import { prettyError } from "@/utils/errors";
+import { useCellActions } from "../../../core/cells/cells";
+import { PythonIcon } from "../cell/code/icons";
+import {
+  getAICompletionBody,
+  mentionsCompletionSource,
+} from "./completion-utils";
+
+const pythonExtensions = [
+  customPythonLanguageSupport(),
+  EditorView.lineWrapping,
+];
+const sqlExtensions = [sql(), EditorView.lineWrapping];
+
+// Persist across sessions
+const languageAtom = atomWithStorage<"python" | "sql">(
+  "marimo:ai-language",
+  "python",
+);
+
+/**
+ * Add a cell with AI.
+ */
+export const AddCellWithAI: React.FC<{
+  onClose: () => void;
+}> = ({ onClose }) => {
+  const { createNewCell } = useCellActions();
+  const [completionBody, setCompletionBody] = useState<object>({});
+  const [language, setLanguage] = useAtom(languageAtom);
+  const { theme } = useTheme();
+  const runtimeManager = useRuntimeManager();
+
+  const ai = useAtomValue(aiAtom);
+  const editModel = ai?.models?.edit_model || DEFAULT_AI_MODEL;
+  const { saveModelChange } = useModelChange();
+
+  const {
+    completion,
+    input,
+    stop,
+    isLoading,
+    setCompletion,
+    setInput,
+    handleSubmit,
+  } = useCompletion({
+    api: runtimeManager.getAiURL("completion").toString(),
+    headers: runtimeManager.headers(),
+    streamProtocol: "text",
+    // Throttle the messages and data updates to 100ms
+    experimental_throttle: 100,
+    body: {
+      ...completionBody,
+      language: language,
+      code: "",
+    },
+    onError: (error) => {
+      toast({
+        title: "Generate with AI failed",
+        description: prettyError(error),
+      });
+    },
+    onFinish: (_prompt, completion) => {
+      // Remove trailing new lines
+      setCompletion(completion.trimEnd());
+    },
+  });
+
+  const submit = () => {
+    if (!isLoading) {
+      handleSubmit();
+    }
+  };
+
+  const pythonIcon = (
+    <>
+      <PythonIcon className="size-4 mr-2" />
+      Python
+    </>
+  );
+
+  const sqlIcon = (
+    <>
+      <DatabaseIcon className="size-4 mr-2" />
+      SQL
+    </>
+  );
+
+  const languageDropdown = (
+    <DropdownMenu modal={false}>
+      <DropdownMenuTrigger asChild={true}>
+        <Button
+          variant="text"
+          className="ml-2"
+          size="xs"
+          data-testid="language-button"
+        >
+          {language === "python" ? pythonIcon : sqlIcon}
+          <ChevronsUpDown className="ml-1 h-3.5 w-3.5 text-muted-foreground/70" />
+        </Button>
+      </DropdownMenuTrigger>
+      <DropdownMenuContent align="center">
+        <div className="px-2 py-1 font-semibold">Select language</div>
+        <DropdownMenuSeparator />
+        <DropdownMenuItem onClick={() => setLanguage("python")}>
+          {pythonIcon}
+        </DropdownMenuItem>
+        <DropdownMenuItem onClick={() => setLanguage("sql")}>
+          {sqlIcon}
+        </DropdownMenuItem>
+      </DropdownMenuContent>
+    </DropdownMenu>
+  );
+
+  const inputComponent = (
+    <div className="flex items-center px-3">
+      <SparklesIcon className="size-4 text-(--blue-11) mr-2" />
+      <PromptInput
+        onClose={() => {
+          setCompletion("");
+          onClose();
+        }}
+        value={input}
+        onChange={(newValue) => {
+          setInput(newValue);
+          setCompletionBody(getAICompletionBody({ input: newValue }));
+        }}
+        onSubmit={submit}
+      />
+      {isLoading && (
+        <Button
+          data-testid="stop-completion-button"
+          variant="text"
+          size="sm"
+          className="mb-0"
+          onClick={stop}
+        >
+          <Loader2Icon className="animate-spin mr-1" size={14} />
+          Stop
+        </Button>
+      )}
+      <Button variant="text" size="sm" onClick={submit} title="Submit">
+        <SendHorizontal className="size-4" />
+      </Button>
+      <Button variant="text" size="sm" className="mb-0 px-1" onClick={onClose}>
+        <XIcon className="size-4" />
+      </Button>
+    </div>
+  );
+
+  return (
+    <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 && (
+          <span className="text-xs text-muted-foreground px-3 flex flex-col gap-1">
+            <span>
+              You can mention{" "}
+              <span className="text-(--cyan-11)">@dataframe</span> or{" "}
+              <span className="text-(--cyan-11)">@sql_table</span> to pull
+              additional context such as column names.
+            </span>
+            <span>Code from other cells is automatically included.</span>
+          </span>
+        )}
+        {completion && (
+          <>
+            <Button
+              data-testid="accept-completion-button"
+              variant="text"
+              size="sm"
+              className="mb-0"
+              onClick={() => {
+                createNewCell({
+                  cellId: "__end__",
+                  before: false,
+                  code:
+                    language === "python"
+                      ? completion
+                      : SQLLanguageAdapter.fromQuery(completion),
+                });
+                setCompletion("");
+                onClose();
+              }}
+            >
+              <span className="text-(--grass-11)">Accept</span>
+            </Button>
+            <Button
+              data-testid="decline-completion-button"
+              variant="text"
+              size="sm"
+              className="mb-0 pl-1"
+              onClick={() => setCompletion("")}
+            >
+              <span className="text-(--red-10)">Reject</span>
+            </Button>
+          </>
+        )}
+        <div className="ml-auto flex items-center gap-1">
+          {languageDropdown}
+          <AIModelDropdown
+            value={editModel}
+            onSelect={(model) => {
+              saveModelChange(model, "edit");
+            }}
+            triggerClassName="h-7 text-xs max-w-64"
+            iconSize="small"
+            forRole="edit"
+          />
+        </div>
+      </div>
+
+      {completion && (
+        <ReactCodeMirror
+          value={completion}
+          className="cm border-t"
+          onChange={setCompletion}
+          theme={theme === "dark" ? "dark" : "light"}
+          extensions={language === "python" ? pythonExtensions : sqlExtensions}
+        />
+      )}
+    </div>
+  );
+};
+
+export interface AdditionalCompletions {
+  triggerCompletionRegex: RegExp;
+  completions: Completion[];
+}
+
+interface PromptInputProps {
+  inputRef?: React.RefObject<ReactCodeMirrorRef | null>;
+  placeholder?: string;
+  value: string;
+  className?: string;
+  onClose: () => void;
+  onChange: (value: string) => void;
+  onSubmit: (e: KeyboardEvent | undefined, value: string) => void;
+  additionalCompletions?: AdditionalCompletions;
+  maxHeight?: string;
+}
+
+/**
+ * CodeMirror-based input for the AI prompt.
+ *
+ * This is just text (no language support), but we use codemirror to get autocomplete
+ * for @dataframe and @sql_table.
+ */
+export const PromptInput = ({
+  value,
+  placeholder,
+  inputRef,
+  className,
+  onChange,
+  onSubmit,
+  onClose,
+  additionalCompletions,
+  maxHeight,
+}: PromptInputProps) => {
+  const handleSubmit = onSubmit;
+  const handleEscape = onClose;
+  const store = useStore();
+  const { theme } = useTheme();
+
+  const additionalCompletionsSource: CompletionSource = useEvent(
+    (context: CompletionContext) => {
+      if (!additionalCompletions) {
+        return null;
+      }
+
+      return mentionsCompletionSource(
+        [additionalCompletions.triggerCompletionRegex],
+        additionalCompletions.completions,
+      )(context);
+    },
+  );
+
+  // Changing extensions can be expensive, so
+  // it is worth making sure this is memoized well.
+  const extensions = useMemo(() => {
+    const markdownLanguage = markdown();
+    return [
+      autocompletion({}),
+      markdownLanguage,
+      resourceExtension(markdownLanguage.language, store),
+      markdownLanguage.language.data.of({
+        autocomplete: additionalCompletionsSource,
+      }),
+      EditorView.lineWrapping,
+      minimalSetup(),
+      Prec.highest(
+        keymap.of([
+          {
+            preventDefault: true,
+            stopPropagation: true,
+            any: (view, event) => {
+              const pressedModOrShift =
+                event.metaKey || event.ctrlKey || event.shiftKey;
+              // If no mod key is pressed, submit
+              if (event.key === "Enter" && !pressedModOrShift) {
+                handleSubmit(event, view.state.doc.toString());
+                event.preventDefault();
+                event.stopPropagation();
+                return true;
+              }
+              // Mod+Enter does add a new line already by codemirror
+              // But Shift+Enter does not, so we need to handle it manually
+              if (event.key === "Enter" && event.shiftKey) {
+                const cursorPosition = view.state.selection.main.from;
+                // Insert a new line
+                view.dispatch({
+                  changes: {
+                    from: cursorPosition,
+                    to: cursorPosition,
+                    insert: "\n",
+                  },
+                  selection: {
+                    anchor: cursorPosition + 1,
+                    head: cursorPosition + 1,
+                  },
+                });
+                event.preventDefault();
+                event.stopPropagation();
+                return true;
+              }
+
+              return false;
+            },
+          },
+        ]),
+      ),
+      keymap.of([
+        {
+          key: "Escape",
+          preventDefault: true,
+          stopPropagation: true,
+          run: () => {
+            handleEscape();
+            return true;
+          },
+        },
+      ]),
+      // Trap arrow up/down to prevent them from being used to navigate the editor
+      keymap.of([
+        {
+          key: "ArrowUp",
+          preventDefault: true,
+          stopPropagation: true,
+        },
+      ]),
+      keymap.of([
+        {
+          key: "ArrowDown",
+          preventDefault: true,
+          stopPropagation: true,
+        },
+      ]),
+    ];
+  }, [store, additionalCompletionsSource, handleSubmit, handleEscape]);
+
+  return (
+    <ReactCodeMirror
+      ref={inputRef}
+      className={cn("flex-1 font-sans overflow-auto my-1", className)}
+      width="100%"
+      maxHeight={maxHeight}
+      value={value}
+      basicSetup={false}
+      extensions={extensions}
+      onChange={onChange}
+      theme={theme === "dark" ? "dark" : "light"}
+      placeholder={placeholder || "Generate with AI"}
+    />
+  );
+};

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

@@ -0,0 +1,261 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { useCompletion } from "@ai-sdk/react";
+import { EditorView } from "@codemirror/view";
+import { Loader2Icon, SparklesIcon, XIcon } from "lucide-react";
+import React, { useEffect, useId, useState } from "react";
+import CodeMirrorMerge from "react-codemirror-merge";
+import { Button } from "@/components/ui/button";
+import { customPythonLanguageSupport } from "@/core/codemirror/language/languages/python";
+
+import "./merge-editor.css";
+import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
+import { useAtom } from "jotai";
+import { Checkbox } from "@/components/ui/checkbox";
+import { Label } from "@/components/ui/label";
+import { Tooltip } from "@/components/ui/tooltip";
+import { toast } from "@/components/ui/use-toast";
+import { includeOtherCellsAtom } from "@/core/ai/state";
+import { getCodes } from "@/core/codemirror/copilot/getCodes";
+import type { LanguageAdapterType } from "@/core/codemirror/language/types";
+import { selectAllText } from "@/core/codemirror/utils";
+import { useRuntimeManager } from "@/core/runtime/config";
+import { useTheme } from "@/theme/useTheme";
+import { cn } from "@/utils/cn";
+import { prettyError } from "@/utils/errors";
+import { retryWithTimeout } from "@/utils/timeout";
+import { PromptInput } from "./add-cell-with-ai";
+import { getAICompletionBody } from "./completion-utils";
+
+const Original = CodeMirrorMerge.Original;
+const Modified = CodeMirrorMerge.Modified;
+
+interface Props {
+  className?: string;
+  currentCode: string;
+  currentLanguageAdapter: LanguageAdapterType | undefined;
+  initialPrompt: string | undefined;
+  onChange: (code: string) => void;
+  declineChange: () => void;
+  acceptChange: (rightHandCode: string) => void;
+  enabled: boolean;
+  /**
+   * Children shown when there is no completion
+   */
+  children: React.ReactNode;
+}
+
+const baseExtensions = [customPythonLanguageSupport(), EditorView.lineWrapping];
+
+/**
+ * Editor for AI completions that goes above a cell to modify it.
+ *
+ * This shows a left/right split with the original and modified code.
+ */
+export const AiCompletionEditor: React.FC<Props> = ({
+  className,
+  onChange,
+  initialPrompt,
+  currentLanguageAdapter,
+  currentCode,
+  declineChange,
+  acceptChange,
+  enabled,
+  children,
+}) => {
+  const [completionBody, setCompletionBody] = useState<object>({});
+
+  const [includeOtherCells, setIncludeOtherCells] = useAtom(
+    includeOtherCellsAtom,
+  );
+  const includeOtherCellsCheckboxId = useId();
+
+  const runtimeManager = useRuntimeManager();
+
+  const {
+    completion: untrimmedCompletion,
+    input,
+    stop,
+    isLoading,
+    setCompletion,
+    setInput,
+    handleSubmit,
+  } = useCompletion({
+    api: runtimeManager.getAiURL("completion").toString(),
+    headers: runtimeManager.headers(),
+    initialInput: initialPrompt,
+    streamProtocol: "text",
+    // Throttle the messages and data updates to 100ms
+    experimental_throttle: 100,
+    body: {
+      ...completionBody,
+      includeOtherCode: includeOtherCells ? getCodes(currentCode) : "",
+      code: currentCode,
+      language: currentLanguageAdapter,
+    },
+    onError: (error) => {
+      toast({
+        title: "Completion failed",
+        description: prettyError(error),
+      });
+    },
+    onFinish: (_prompt, completion) => {
+      // Remove trailing new lines
+      setCompletion(completion.trimEnd());
+    },
+  });
+
+  const inputRef = React.useRef<ReactCodeMirrorRef>(null);
+  const completion = untrimmedCompletion.trimEnd();
+
+  // Focus the input
+  useEffect(() => {
+    if (enabled) {
+      retryWithTimeout(
+        () => {
+          const input = inputRef.current;
+          if (input?.view) {
+            input.view.focus();
+            return true;
+          }
+          return false;
+        },
+        { retries: 3, delay: 100, initialDelay: 100 },
+      ); // Wait for animation to complete
+
+      selectAllText(inputRef.current?.view);
+    }
+  }, [enabled]);
+
+  // Reset the input when the prompt changes
+  useEffect(() => {
+    if (enabled) {
+      setInput(initialPrompt || "");
+    }
+  }, [enabled, initialPrompt, setInput]);
+
+  const { theme } = useTheme();
+
+  return (
+    <div className={cn("flex flex-col w-full rounded-[inherit]", className)}>
+      <div
+        className={cn(
+          "flex items-center gap-2 border-b px-3 transition-all rounded-[inherit] rounded-b-none duration-300",
+          enabled && "max-h-[400px] min-h-11 visible",
+          !enabled && "max-h-0 min-h-0 invisible",
+        )}
+      >
+        {enabled && (
+          <>
+            <SparklesIcon className="text-(--blue-10) shrink-0" size={16} />
+            <PromptInput
+              inputRef={inputRef}
+              onClose={() => {
+                declineChange();
+                setCompletion("");
+              }}
+              value={input}
+              onChange={(newValue) => {
+                setInput(newValue);
+                setCompletionBody(getAICompletionBody({ input: newValue }));
+              }}
+              onSubmit={() => {
+                if (!isLoading) {
+                  handleSubmit();
+                }
+              }}
+            />
+            {isLoading && (
+              <Button
+                data-testid="stop-completion-button"
+                variant="text"
+                size="xs"
+                className="mb-0"
+                onClick={stop}
+              >
+                <Loader2Icon className="animate-spin mr-1" size={14} />
+                Stop
+              </Button>
+            )}
+            {completion && (
+              <>
+                <Button
+                  data-testid="accept-completion-button"
+                  variant="text"
+                  size="xs"
+                  className="mb-0"
+                  disabled={isLoading}
+                  onClick={() => {
+                    acceptChange(completion);
+                    setCompletion("");
+                  }}
+                >
+                  <span className="text-(--grass-11) opacity-100">Accept</span>
+                </Button>
+                <Button
+                  data-testid="decline-completion-button"
+                  variant="text"
+                  size="xs"
+                  className="mb-0 pl-1"
+                  onClick={() => {
+                    setCompletion("");
+                  }}
+                >
+                  <span className="text-(--red-10)">Reject</span>
+                </Button>
+              </>
+            )}
+            <div className="h-full w-px bg-border mx-2" />
+            <Tooltip content="Include code from other cells">
+              <div className="flex flex-row items-start gap-1 overflow-hidden">
+                <Checkbox
+                  data-testid="include-other-cells-checkbox"
+                  id={includeOtherCellsCheckboxId}
+                  checked={includeOtherCells}
+                  onCheckedChange={(checked) =>
+                    setIncludeOtherCells(Boolean(checked))
+                  }
+                />
+                <Label
+                  htmlFor={includeOtherCellsCheckboxId}
+                  className="text-muted-foreground text-xs whitespace-nowrap ellipsis"
+                >
+                  Include all code
+                </Label>
+              </div>
+            </Tooltip>
+            <Button
+              data-testid="decline-completion-button"
+              variant="text"
+              size="icon"
+              disabled={isLoading}
+              onClick={() => {
+                stop();
+                declineChange();
+                setCompletion("");
+              }}
+            >
+              <XIcon className="text-(--red-10)" size={16} />
+            </Button>
+          </>
+        )}
+      </div>
+      {completion && enabled && (
+        <CodeMirrorMerge className="cm" theme={theme}>
+          <Original
+            onChange={onChange}
+            value={currentCode}
+            extensions={baseExtensions}
+          />
+          <Modified
+            value={completion}
+            editable={false}
+            readOnly={true}
+            extensions={baseExtensions}
+          />
+        </CodeMirrorMerge>
+      )}
+      {(!completion || !enabled) && children}
+    </div>
+  );
+};