fix: better chatbot input & bubble

frontend/package.json

@@ -79,7 +79,6 @@
     "@types/react-grid-layout": "^1.3.5",
     "@uidotdev/usehooks": "^2.4.1",
     "@uiw/codemirror-extensions-langs": "^4.23.5",
-    "@uiw/codemirror-extensions-mentions": "^4.23.5",
     "@uiw/react-codemirror": "^4.23.5",
     "@valtown/codemirror-codeium": "^1.1.1",
     "@xterm/addon-attach": "^0.11.0",

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

@@ -16,7 +16,6 @@ import ReactCodeMirror, {
 import { Prec } from "@codemirror/state";
 import { customPythonLanguageSupport } from "@/core/codemirror/language/python";
 import { asURL } from "@/utils/url";
-import { mentions } from "@uiw/codemirror-extensions-mentions";
 import { useMemo, useState } from "react";
 import { datasetTablesAtom } from "@/core/datasets/state";
 import { useAtom, useAtomValue } from "jotai";
@@ -31,7 +30,7 @@ import { sql } from "@codemirror/lang-sql";
 import { SQLLanguageAdapter } from "@/core/codemirror/language/sql";
 import { atomWithStorage } from "jotai/utils";
 import { type ResolvedTheme, useTheme } from "@/theme/useTheme";
-import { getAICompletionBody } from "./completion-utils";
+import { getAICompletionBody, mentions } from "./completion-utils";
 
 const pythonExtensions = [
   customPythonLanguageSupport(),
@@ -190,6 +189,11 @@ export const AddCellWithAI: React.FC<{
   );
 };
 
+export interface AdditionalCompletions {
+  triggerCompletionRegex: RegExp;
+  completions: Completion[];
+}
+
 interface PromptInputProps {
   inputRef?: React.RefObject<ReactCodeMirrorRef>;
   placeholder?: string;
@@ -198,7 +202,9 @@ interface PromptInputProps {
   onClose: () => void;
   onChange: (value: string) => void;
   onSubmit: (e: KeyboardEvent | undefined, value: string) => void;
+  additionalCompletions?: AdditionalCompletions;
   theme: ResolvedTheme;
+  maxHeight?: string;
 }
 
 /**
@@ -215,7 +221,9 @@ export const PromptInput = ({
   onChange,
   onSubmit,
   onClose,
+  additionalCompletions,
   theme,
+  maxHeight,
 }: PromptInputProps) => {
   const handleSubmit = onSubmit;
   const handleEscape = onClose;
@@ -277,8 +285,16 @@ export const PromptInput = ({
       }),
     );
 
+    const matchBeforeRegexes = [/@(\w+)?/]; // Trigger autocompletion for text that begins with @
+    if (additionalCompletions) {
+      matchBeforeRegexes.push(additionalCompletions.triggerCompletionRegex);
+    }
+    const allCompletions = additionalCompletions
+      ? [...completions, ...additionalCompletions.completions]
+      : completions;
+
     return [
-      mentions(completions),
+      mentions(matchBeforeRegexes, allCompletions),
       EditorView.lineWrapping,
       minimalSetup(),
       Prec.highest(
@@ -349,14 +365,15 @@ export const PromptInput = ({
         },
       ]),
     ];
-  }, [tables, handleSubmit, handleEscape]);
+  }, [tables, additionalCompletions, handleSubmit, handleEscape]);
 
   return (
     <ReactCodeMirror
       ref={inputRef}
       className={cn("flex-1 font-sans overflow-auto my-1", className)}
       autoFocus={true}
       width="100%"
+      maxHeight={maxHeight}
       value={value}
       basicSetup={false}
       extensions={extensions}

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

@@ -6,6 +6,12 @@ import type { AiCompletionRequest } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
 import { Logger } from "@/utils/Logger";
 import { Maps } from "@/utils/maps";
+import {
+  autocompletion,
+  type Completion,
+  type CompletionContext,
+} from "@codemirror/autocomplete";
+import type { Extension } from "@codemirror/state";
 
 /**
  * Gets the request body for the AI completion API.
@@ -48,3 +54,32 @@ function extractDatasets(input: string): DataTable[] {
     .map((name) => existingDatasets.get(name))
     .filter(Boolean);
 }
+
+/**
+ * Adapted from @uiw/codemirror-extensions-mentions
+ * Allows you to specify a custom regex to trigger the autocompletion.
+ */
+export function mentions(
+  matchBeforeRegexes: RegExp[],
+  data: Completion[] = [],
+): Extension {
+  return autocompletion({
+    override: [
+      (context: CompletionContext) => {
+        const word = matchBeforeRegexes
+          .map((regex) => context.matchBefore(regex))
+          .find(Boolean);
+        if (!word) {
+          return null;
+        }
+        if (word && word.from === word.to && !context.explicit) {
+          return null;
+        }
+        return {
+          from: word?.from,
+          options: [...data],
+        };
+      },
+    ],
+  });
+}

frontend/src/core/codemirror/utils.ts

@@ -25,6 +25,18 @@ export function isAtEndOfEditor(ev: { state: EditorState }, isVim = false) {
   return main.from === docLength && main.to === docLength;
 }
 
+export function moveToEndOfEditor(ev: EditorView | undefined) {
+  if (!ev) {
+    return;
+  }
+  ev.dispatch({
+    selection: {
+      anchor: ev.state.doc.length,
+      head: ev.state.doc.length,
+    },
+  });
+}
+
 export function isInVimNormalMode(ev: EditorView): boolean {
   const vimState = getCM(ev)?.state.vim;
   if (!vimState) {

frontend/src/plugins/impl/chat/chat-ui.tsx

@@ -47,6 +47,13 @@ import { renderHTML } from "@/plugins/core/RenderHTML";
 import { Input } from "@/components/ui/input";
 import { PopoverAnchor } from "@radix-ui/react-popover";
 import { copyToClipboard } from "@/utils/copy";
+import {
+  type AdditionalCompletions,
+  PromptInput,
+} from "@/components/editor/ai/add-cell-with-ai";
+import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
+import { useTheme } from "@/theme/useTheme";
+import { moveToEndOfEditor } from "@/core/codemirror/utils";
 
 interface Props {
   prompts: string[];
@@ -60,18 +67,19 @@ interface Props {
 }
 
 export const Chatbot: React.FC<Props> = (props) => {
-  const inputRef = React.useRef<HTMLInputElement>(null);
   const [config, setConfig] = useState<ChatConfig>(props.config);
   const [files, setFiles] = useState<FileList | undefined>(undefined);
   const fileInputRef = useRef<HTMLInputElement>(null);
+  const formRef = useRef<HTMLFormElement>(null);
+  const codeMirrorInputRef = useRef<ReactCodeMirrorRef>(null);
   const scrollContainerRef = useRef<HTMLDivElement>(null);
+  const { theme } = useTheme();
 
   const {
     messages,
     setMessages,
     input,
     setInput,
-    handleInputChange,
     handleSubmit,
     isLoading,
     stop,
@@ -196,6 +204,20 @@ export const Chatbot: React.FC<Props> = (props) => {
       props.allowAttachments.length > 0) ||
     props.allowAttachments === true;
 
+  const promptCompletions: AdditionalCompletions = {
+    // sentence has to begin with '/' to trigger autocomplete
+    triggerCompletionRegex: /^\/(\w+)?/,
+    completions: props.prompts.map((prompt) => ({
+      label: `/${prompt}`,
+      displayLabel: prompt,
+      apply: prompt,
+    })),
+  };
+  const promptInputPlaceholder =
+    props.prompts.length > 0
+      ? "Type your message here, / for prompts"
+      : "Type your message here...";
+
   useEffect(() => {
     // When the message length changes, scroll to the bottom
     scrollContainerRef.current?.scrollTo({
@@ -242,7 +264,11 @@ export const Chatbot: React.FC<Props> = (props) => {
                   : "bg-[var(--slate-4)] text-[var(--slate-12)]"
               }`}
             >
-              <p>{renderMessage(message)}</p>
+              <p
+                className={cn(message.role === "user" && "whitespace-pre-wrap")}
+              >
+                {renderMessage(message)}
+              </p>
             </div>
             <div className="flex justify-end text-xs gap-2 invisible group-hover:visible">
               <button
@@ -298,6 +324,7 @@ export const Chatbot: React.FC<Props> = (props) => {
             experimental_attachments: files,
           });
         }}
+        ref={formRef}
         className="flex w-full border-t border-[var(--slate-6)] px-2 py-1 items-center"
       >
         {props.showConfigurationControls && (
@@ -309,22 +336,30 @@ export const Chatbot: React.FC<Props> = (props) => {
             onSelect={(prompt) => {
               setInput(prompt);
               requestAnimationFrame(() => {
-                inputRef.current?.focus();
-                inputRef.current?.setSelectionRange(
-                  prompt.length,
-                  prompt.length,
-                );
+                codeMirrorInputRef.current?.view?.focus();
+                moveToEndOfEditor(codeMirrorInputRef.current?.view);
               });
             }}
           />
         )}
-        <input
-          name="prompt"
-          ref={inputRef}
+        <PromptInput
+          className="rounded-sm mr-2"
+          placeholder={promptInputPlaceholder}
           value={input}
-          onChange={handleInputChange}
-          className="flex w-full outline-none bg-transparent ml-2 text-[var(--slate-12)] mr-2"
-          placeholder="Type your message..."
+          inputRef={codeMirrorInputRef}
+          theme={theme}
+          maxHeight={props.maxHeight ? `${props.maxHeight / 2}px` : undefined}
+          onChange={setInput}
+          onSubmit={(_evt, newValue) => {
+            if (!newValue.trim()) {
+              return;
+            }
+            formRef.current?.requestSubmit();
+          }}
+          onClose={() => {
+            // no-op
+          }}
+          additionalCompletions={promptCompletions}
         />
         {files && files.length === 1 && (
           <span