From b01da38eaaa5346cb81df4e49641d9cfea4960b4 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Sun, 12 Jan 2025 19:39:20 +0800
Subject: [PATCH 1/6] render chat bubble w/o breaking formatting

---
 frontend/src/plugins/impl/chat/chat-ui.tsx | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/frontend/src/plugins/impl/chat/chat-ui.tsx b/frontend/src/plugins/impl/chat/chat-ui.tsx
index d413af708da..bdf010f298b 100644
--- a/frontend/src/plugins/impl/chat/chat-ui.tsx
+++ b/frontend/src/plugins/impl/chat/chat-ui.tsx
@@ -242,7 +242,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

From 0e6ff833a0fb7eb11bb2843da074ce7a3ab65a22 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Mon, 13 Jan 2025 01:16:14 +0800
Subject: [PATCH 2/6] use codemirror as input

---
 .../components/editor/ai/add-cell-with-ai.tsx | 23 ++++++--
 .../components/editor/ai/completion-utils.ts  | 33 +++++++++++
 frontend/src/core/codemirror/utils.ts         | 12 ++++
 frontend/src/plugins/impl/chat/chat-ui.tsx    | 56 ++++++++++++++-----
 4 files changed, 107 insertions(+), 17 deletions(-)

diff --git a/frontend/src/components/editor/ai/add-cell-with-ai.tsx b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
index f93a16fcf52..a0295212406 100644
--- a/frontend/src/components/editor/ai/add-cell-with-ai.tsx
+++ b/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 {
+  triggerSymbol: string; // Text that begins with triggerSymbol will trigger autocompletion
+  completions: Completion[];
+}
+
 interface PromptInputProps {
   inputRef?: React.RefObject<ReactCodeMirrorRef>;
   placeholder?: string;
@@ -198,6 +202,7 @@ interface PromptInputProps {
   onClose: () => void;
   onChange: (value: string) => void;
   onSubmit: (e: KeyboardEvent | undefined, value: string) => void;
+  additionalCompletions?: AdditionalCompletions;
   theme: ResolvedTheme;
 }
 
@@ -215,6 +220,7 @@ export const PromptInput = ({
   onChange,
   onSubmit,
   onClose,
+  additionalCompletions,
   theme,
 }: PromptInputProps) => {
   const handleSubmit = onSubmit;
@@ -277,8 +283,17 @@ export const PromptInput = ({
       }),
     );
 
+    // Trigger autocompletion for text that begins with @ or
+    // @ + additional symbols specified
+    const matchBeforeRegex = additionalCompletions
+      ? new RegExp(`[@${additionalCompletions.triggerSymbol}](\\w+)?`)
+      : /@(\w+)?/;
+    const allCompletions = additionalCompletions
+      ? [...completions, ...additionalCompletions.completions]
+      : completions;
+
     return [
-      mentions(completions),
+      mentions(matchBeforeRegex, allCompletions),
       EditorView.lineWrapping,
       minimalSetup(),
       Prec.highest(
@@ -349,7 +364,7 @@ export const PromptInput = ({
         },
       ]),
     ];
-  }, [tables, handleSubmit, handleEscape]);
+  }, [tables, additionalCompletions, handleSubmit, handleEscape]);
 
   return (
     <ReactCodeMirror
diff --git a/frontend/src/components/editor/ai/completion-utils.ts b/frontend/src/components/editor/ai/completion-utils.ts
index eab20dbc5a6..668ec1679a7 100644
--- a/frontend/src/components/editor/ai/completion-utils.ts
+++ b/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,30 @@ 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(
+  matchBeforeRegex: RegExp,
+  data: Completion[] = [],
+): Extension {
+  return autocompletion({
+    override: [
+      (context: CompletionContext) => {
+        const word = context.matchBefore(matchBeforeRegex);
+        if (!word) {
+          return null;
+        }
+        if (word && word.from === word.to && !context.explicit) {
+          return null;
+        }
+        return {
+          from: word?.from,
+          options: [...data],
+        };
+      },
+    ],
+  });
+}
diff --git a/frontend/src/core/codemirror/utils.ts b/frontend/src/core/codemirror/utils.ts
index f4c05efe224..fb769bc6512 100644
--- a/frontend/src/core/codemirror/utils.ts
+++ b/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) {
diff --git a/frontend/src/plugins/impl/chat/chat-ui.tsx b/frontend/src/plugins/impl/chat/chat-ui.tsx
index bdf010f298b..a8f07a9b8fe 100644
--- a/frontend/src/plugins/impl/chat/chat-ui.tsx
+++ b/frontend/src/plugins/impl/chat/chat-ui.tsx
@@ -47,6 +47,14 @@ 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";
+import type { Completion } from "@codemirror/autocomplete";
 
 interface Props {
   prompts: string[];
@@ -60,18 +68,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 +205,17 @@ export const Chatbot: React.FC<Props> = (props) => {
       props.allowAttachments.length > 0) ||
     props.allowAttachments === true;
 
+  const promptCompletions: AdditionalCompletions = {
+    triggerSymbol: "/",
+    completions: props.prompts.map(
+      (prompt): Completion => ({
+        label: `/${prompt}`,
+        displayLabel: prompt,
+        apply: prompt,
+      }),
+    ),
+  };
+
   useEffect(() => {
     // When the message length changes, scroll to the bottom
     scrollContainerRef.current?.scrollTo({
@@ -302,6 +322,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 && (
@@ -313,22 +334,31 @@ 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"
+          placeholder="Type your message here, / for prompts"
           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}
+          onChange={setInput}
+          onSubmit={(_evt, newValue) => {
+            if (!newValue.trim()) {
+              return;
+            }
+            if (formRef.current) {
+              formRef.current.requestSubmit();
+            }
+          }}
+          onClose={() => {
+            // no-op
+          }}
+          additionalCompletions={promptCompletions}
         />
         {files && files.length === 1 && (
           <span

From 4f72fdd088ab160fb299ddcb50d77c3a52e5bba7 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Mon, 13 Jan 2025 01:18:41 +0800
Subject: [PATCH 3/6] small refactor

---
 frontend/src/plugins/impl/chat/chat-ui.tsx | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/frontend/src/plugins/impl/chat/chat-ui.tsx b/frontend/src/plugins/impl/chat/chat-ui.tsx
index a8f07a9b8fe..fd8359f3877 100644
--- a/frontend/src/plugins/impl/chat/chat-ui.tsx
+++ b/frontend/src/plugins/impl/chat/chat-ui.tsx
@@ -351,9 +351,7 @@ export const Chatbot: React.FC<Props> = (props) => {
             if (!newValue.trim()) {
               return;
             }
-            if (formRef.current) {
-              formRef.current.requestSubmit();
-            }
+            formRef.current?.requestSubmit();
           }}
           onClose={() => {
             // no-op

From d68f247f203e8b00fb0507c8ef9f5e6fdf4103ea Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Mon, 13 Jan 2025 01:20:32 +0800
Subject: [PATCH 4/6] remove unecc extension

---
 frontend/package.json   |  1 -
 frontend/pnpm-lock.yaml | 14 --------------
 2 files changed, 15 deletions(-)

diff --git a/frontend/package.json b/frontend/package.json
index addaaa7536b..abe675cf890 100644
--- a/frontend/package.json
+++ b/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",
diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml
index af2a5228560..794293ff559 100644
--- a/frontend/pnpm-lock.yaml
+++ b/frontend/pnpm-lock.yaml
@@ -213,9 +213,6 @@ importers:
       '@uiw/codemirror-extensions-langs':
         specifier: ^4.23.5
         version: 4.23.5(@codemirror/autocomplete@6.18.3(@codemirror/language@6.10.3)(@codemirror/state@6.4.1)(@codemirror/view@6.34.3)(@lezer/common@1.2.1))(@codemirror/language-data@6.5.1(@codemirror/view@6.34.3))(@codemirror/language@6.10.3)(@codemirror/legacy-modes@6.4.1)(@codemirror/state@6.4.1)(@codemirror/view@6.34.3)(@lezer/common@1.2.1)(@lezer/highlight@1.2.1)(@lezer/javascript@1.4.18)(@lezer/lr@1.4.2)
-      '@uiw/codemirror-extensions-mentions':
-        specifier: ^4.23.5
-        version: 4.23.5(@codemirror/state@6.4.1)(@codemirror/view@6.34.3)
       '@uiw/react-codemirror':
         specifier: ^4.23.5
         version: 4.23.5(@babel/runtime@7.25.6)(@codemirror/autocomplete@6.18.3(@codemirror/language@6.10.3)(@codemirror/state@6.4.1)(@codemirror/view@6.34.3)(@lezer/common@1.2.1))(@codemirror/language@6.10.3)(@codemirror/lint@6.8.2)(@codemirror/search@6.5.7)(@codemirror/state@6.4.1)(@codemirror/theme-one-dark@6.1.2)(@codemirror/view@6.34.3)(codemirror@6.0.1(@lezer/common@1.2.1))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)
@@ -3654,12 +3651,6 @@ packages:
       '@codemirror/language-data': '>=6.0.0'
       '@codemirror/legacy-modes': '>=6.0.0'
 
-  '@uiw/codemirror-extensions-mentions@4.23.5':
-    resolution: {integrity: sha512-ZxX/KXY9qmgb32u4Py/+MYqKL9ln13zJBjr/G/GFDjyUBGA4I5XWv4yQczBdXg+kyAyiP2tqWURoQfvSGtV12g==}
-    peerDependencies:
-      '@codemirror/state': '>=6.0.0'
-      '@codemirror/view': '>=6.0.0'
-
   '@uiw/react-codemirror@4.23.5':
     resolution: {integrity: sha512-2zzGpx61L4mq9zDG/hfsO4wAH209TBE8VVsoj/qrccRe6KfcneCwKgRxtQjxBCCnO0Q5S+IP+uwCx5bXRzgQFQ==}
     peerDependencies:
@@ -12906,11 +12897,6 @@ snapshots:
       - '@lezer/javascript'
       - '@lezer/lr'
 
-  '@uiw/codemirror-extensions-mentions@4.23.5(@codemirror/state@6.4.1)(@codemirror/view@6.34.3)':
-    dependencies:
-      '@codemirror/state': 6.4.1
-      '@codemirror/view': 6.34.3
-
   '@uiw/react-codemirror@4.23.5(@babel/runtime@7.25.6)(@codemirror/autocomplete@6.18.3(@codemirror/language@6.10.3)(@codemirror/state@6.4.1)(@codemirror/view@6.34.3)(@lezer/common@1.2.1))(@codemirror/language@6.10.3)(@codemirror/lint@6.8.2)(@codemirror/search@6.5.7)(@codemirror/state@6.4.1)(@codemirror/theme-one-dark@6.1.2)(@codemirror/view@6.34.3)(codemirror@6.0.1(@lezer/common@1.2.1))(react-dom@18.3.1(react@18.3.1))(react@18.3.1)':
     dependencies:
       '@babel/runtime': 7.25.6

From 4b38701fb2de272f82cb2090db5b2e2e488d5131 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Mon, 13 Jan 2025 01:35:29 +0800
Subject: [PATCH 5/6] modify comment

---
 frontend/src/components/editor/ai/add-cell-with-ai.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/frontend/src/components/editor/ai/add-cell-with-ai.tsx b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
index a0295212406..bdd9740c85b 100644
--- a/frontend/src/components/editor/ai/add-cell-with-ai.tsx
+++ b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
@@ -190,7 +190,7 @@ export const AddCellWithAI: React.FC<{
 };
 
 export interface AdditionalCompletions {
-  triggerSymbol: string; // Text that begins with triggerSymbol will trigger autocompletion
+  triggerSymbol: string; // Symbol that will trigger autocompletion when text begins with it
   completions: Completion[];
 }
 

From dc0803d725208279b2923f0df41b0c2f376638ce Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Mon, 13 Jan 2025 01:42:30 +0800
Subject: [PATCH 6/6] small right margin

---
 frontend/src/plugins/impl/chat/chat-ui.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/frontend/src/plugins/impl/chat/chat-ui.tsx b/frontend/src/plugins/impl/chat/chat-ui.tsx
index fd8359f3877..c8765d50e79 100644
--- a/frontend/src/plugins/impl/chat/chat-ui.tsx
+++ b/frontend/src/plugins/impl/chat/chat-ui.tsx
@@ -341,7 +341,7 @@ export const Chatbot: React.FC<Props> = (props) => {
           />
         )}
         <PromptInput
-          className="rounded-sm"
+          className="rounded-sm mr-2"
           placeholder="Type your message here, / for prompts"
           value={input}
           inputRef={codeMirrorInputRef}