Feat(notetaker): hide selected text from searchbar + handle edge cases for AI selection context

frontend/lib/components/TextEditor/TextEditor.tsx

@@ -2,7 +2,7 @@
 import { Editor, Extension } from "@tiptap/core";
 import Focus from "@tiptap/extension-focus";
 import { Link } from "@tiptap/extension-link";
-import { BubbleMenu, EditorContent, useEditor } from "@tiptap/react";
+import { EditorContent, useEditor } from "@tiptap/react";
 import { StarterKit } from "@tiptap/starter-kit";
 import { useRef, useState } from "react";
 
@@ -12,8 +12,7 @@ import styles from "./TextEditor.module.scss";
 import { TextEditorSearchBar } from "./components/TextEditorSearchBar/TextEditorSearchBar";
 import { Toolbar } from "./components/Toolbar/Toolbar";
 import { AIHighlight } from "./extensions/AIHighlight";
-
-import { QuivrButton } from "../ui/QuivrButton/QuivrButton";
+import { AiResponse } from "./extensions/AiResponse";
 
 const defaultContent = `
   <h1>My Note</h1>
@@ -35,21 +34,17 @@ export const TextEditor = (): JSX.Element => {
     addKeyboardShortcuts: () => {
       return {
         "Mod-f": ({ editor }) => {
-          const selection = editor.state.doc.textBetween(
+          const content = editor.state.doc.textBetween(
             editor.state.selection.from,
             editor.state.selection.to
           );
 
-          if (selection) {
+          if (content) {
             editor.commands.setSelectionHighlight();
           }
 
           setSearchBarOpen(true);
-          searchEditorRef.current
-            ?.chain()
-            .focus()
-            .setContent(selection, undefined, { preserveWhitespace: true })
-            .run();
+          searchEditorRef.current?.chain().focus().run();
 
           return true;
         },
@@ -75,6 +70,7 @@ export const TextEditor = (): JSX.Element => {
             class: styles.ai_highlight,
           },
         }),
+        AiResponse,
       ],
       content: defaultContent,
       immediatelyRender: false,
@@ -95,42 +91,11 @@ export const TextEditor = (): JSX.Element => {
   return (
     <div className={styles.main_container}>
       <div className={styles.editor_wrapper}>
-        <Toolbar toggleSearchBar={toggleSearchBar} editor={editor} />
-        <div>
-          <BubbleMenu
-            shouldShow={() => {
-              return editor.isActive(AIHighlight.name, { type: "ai" });
-            }}
-            tippyOptions={{
-              moveTransition: "transform 0.1s",
-              placement: "bottom-start",
-            }}
-            className={styles.bubble_menu}
-            editor={editor}
-          >
-            <QuivrButton
-              onClick={() => {
-                editor
-                  .chain()
-                  .extendMarkRange(AIHighlight.name)
-                  .unsetHighlight()
-                  .run();
-              }}
-              label="Accept"
-              color="primary"
-              iconName="check"
-            />
-            <QuivrButton
-              onClick={() => {
-                editor.commands.undo();
-              }}
-              label="Decline"
-              color="dangerous"
-              iconName="close"
-            />
-          </BubbleMenu>
-        </div>
-
+        <Toolbar
+          searchBarOpen={searchBarOpen}
+          toggleSearchBar={toggleSearchBar}
+          editor={editor}
+        />
         <EditorContent className={styles.content_wrapper} editor={editor} />
       </div>
       <div

frontend/lib/components/TextEditor/components/AIResponseNode/AIResponseNode.module.scss

@@ -0,0 +1,25 @@
+@use "styles/Spacings.module.scss";
+@use "styles/Radius.module.scss";
+
+.wrapper {
+  padding: Spacings.$spacing05;
+  display: flex;
+  flex-direction: column;
+  gap: Spacings.$spacing07;
+  border: solid 1px var(--border-2);
+  background-color: var(--background-1);
+  border-radius: Radius.$normal;
+
+  .context {
+    padding: Spacings.$spacing05;
+    border: solid 1px var(--border-2);
+    background-color: var(--background-2);
+    border-radius: Radius.$normal;
+    opacity: 0.5;
+  }
+
+  .actions_wrapper {
+    display: flex;
+    gap: Spacings.$spacing04;
+  }
+}

frontend/lib/components/TextEditor/components/AIResponseNode/AIResponseNode.tsx

@@ -0,0 +1,42 @@
+import { NodeViewProps } from "@tiptap/core";
+import { NodeViewContent, NodeViewWrapper } from "@tiptap/react";
+
+import styles from "./AIResponseNode.module.scss";
+
+import { QuivrButton } from "../../../ui/QuivrButton/QuivrButton";
+
+const AIResponseNode = (props: NodeViewProps): JSX.Element => {
+  const { attrs, content } = props.node;
+  const { editor } = props;
+
+  return (
+    <NodeViewWrapper className={styles.wrapper}>
+      {attrs.context && <div className={styles.context}>{attrs.context}</div>}
+      <NodeViewContent className={styles.content} />
+      <div className={styles.actions_wrapper}>
+        <QuivrButton
+          color="primary"
+          iconName="check"
+          label="Accept"
+          onClick={() => {
+            editor.commands.acceptAiResponse({
+              content: content.textBetween(0, content.size),
+            });
+          }}
+        />
+        <QuivrButton
+          color="dangerous"
+          iconName="close"
+          label="Decline"
+          onClick={() => {
+            editor.commands.declineAiResponse({
+              prevContent: attrs.context as string,
+            });
+          }}
+        />
+      </div>
+    </NodeViewWrapper>
+  );
+};
+
+export default AIResponseNode;

frontend/lib/components/TextEditor/components/TextEditorSearchBar/TextEditorSearchBar.tsx

@@ -26,12 +26,26 @@ const TextEditorSearchBar = forwardRef<Editor, TextEditorSearchBarProps>(
         return;
       }
 
-      editor
-        .chain()
-        .setAiHighlight()
-        .insertContent(messages[0].assistant)
-        .focus()
-        .run();
+      const context = editor.state.doc.textBetween(
+        editor.state.selection.from,
+        editor.state.selection.to
+      );
+
+      if (context) {
+        editor
+          .chain()
+          .deleteSelection()
+          .createAiResponse({
+            content: messages[0].assistant,
+            context: context,
+          })
+          .run();
+      } else {
+        editor.commands.createAiResponse({
+          content: messages[0].assistant,
+          context: null,
+        });
+      }
     }, [messages.length, router, editor, chatInput.generatingAnswer]);
 
     useEffect(() => {
@@ -48,7 +62,15 @@ const TextEditorSearchBar = forwardRef<Editor, TextEditorSearchBarProps>(
         ref={ref}
         onSearch={onSearch}
         newBrain={newBrain}
-        submitQuestion={(question) => submitQuestion(question, false)}
+        submitQuestion={(question) =>
+          submitQuestion(
+            `${editor.state.doc.textBetween(
+              editor.state.selection.from,
+              editor.state.selection.to
+            )} \n ${question}`,
+            false
+          )
+        }
         {...chatInput}
       />
     );

frontend/lib/components/TextEditor/components/Toolbar/Toolbar.tsx

@@ -23,11 +23,13 @@ export const ToolbarSectionSeparator = (): JSX.Element => {
 
 type ToolbarProps = {
   editor: Editor;
+  searchBarOpen: boolean;
   toggleSearchBar: () => void;
 };
 
 export const Toolbar = ({
   editor,
+  searchBarOpen,
   toggleSearchBar,
 }: ToolbarProps): JSX.Element => {
   const [linkModalOpen, setLinkModalOpen] = useState(false);
@@ -180,7 +182,7 @@ export const Toolbar = ({
       <div className={styles.focusSearchBarBtn}>
         <QuivrButton
           onClick={toggleSearchBar}
-          label="Ask Brain"
+          label={searchBarOpen ? "Close Search Bar" : "Ask Brain"}
           color="primary"
           iconName={"chat"}
         />

frontend/lib/components/TextEditor/extensions/AIHighlight.ts

@@ -1,8 +1,6 @@
 import { Command } from "@tiptap/core";
 import { Highlight } from "@tiptap/extension-highlight";
 
-import { AIHighlightOptions } from "./types";
-
 const unsetSelectionsInDocument: Command = ({ dispatch, state }) => {
   const { tr } = state;
   if (!dispatch) {
@@ -12,10 +10,7 @@ const unsetSelectionsInDocument: Command = ({ dispatch, state }) => {
   state.doc.descendants((node, pos) => {
     if (node.isText) {
       node.marks.map((mark) => {
-        if (
-          mark.type.name === AIHighlight.name &&
-          mark.attrs["type"] === "selection"
-        ) {
+        if (mark.type.name === AIHighlight.name) {
           tr.removeMark(pos, pos + node.nodeSize, mark.type);
         }
       });
@@ -27,34 +22,24 @@ const unsetSelectionsInDocument: Command = ({ dispatch, state }) => {
   return true;
 };
 
-export const AIHighlight = Highlight.extend<AIHighlightOptions>({
+export const AIHighlight = Highlight.extend({
   name: "aiHighlight",
-  addAttributes() {
-    return {
-      ...this.parent?.(),
-      type: "selection",
-    };
-  },
+
   addCommands() {
     return {
       ...this.parent?.(),
       unsetSelectionsInDocument: () => unsetSelectionsInDocument,
       setSelectionHighlight: () => {
         return ({ commands }) => {
-          return commands.setHighlight({ type: "selection" });
-        };
-      },
-      setAiHighlight: () => {
-        return ({ commands }) => {
-          return commands.setHighlight({ type: "ai" });
+          return commands.setHighlight();
         };
       },
     };
   },
+
   onSelectionUpdate() {
     this.editor.commands.unsetSelectionsInDocument();
   },
 }).configure({
   multicolor: true,
-  type: "selection",
 });

frontend/lib/components/TextEditor/extensions/AiResponse.tsx

@@ -0,0 +1,92 @@
+import { mergeAttributes, Node } from "@tiptap/core";
+import { ReactNodeViewRenderer } from "@tiptap/react";
+
+import { AIResponseOptions } from "./types";
+
+import AIResponseNode from "../components/AIResponseNode/AIResponseNode";
+
+export const AiResponse = Node.create<AIResponseOptions>({
+  name: "aiResponse",
+  group: "block",
+  content: "inline*",
+
+  addAttributes() {
+    return {
+      ...this.parent?.(),
+      context: {
+        default: "",
+      },
+    };
+  },
+
+  renderHTML: ({ HTMLAttributes }) => {
+    return ["ai-response", mergeAttributes(HTMLAttributes)];
+  },
+
+  parseHTML: () => {
+    return [
+      {
+        tag: "ai-response",
+      },
+    ];
+  },
+
+  addCommands() {
+    return {
+      ...this.parent?.(),
+      createAiResponse:
+        ({ content, context }: { content: string; context: string }) =>
+        ({ chain }) => {
+          return chain()
+            .insertContent({
+              type: this.name,
+              attrs: { context },
+              content: [
+                {
+                  type: "text",
+                  text: content,
+                },
+              ],
+            })
+            .focus()
+            .run();
+        },
+      acceptAiResponse:
+        ({ content }: { content: string }) =>
+        ({ chain }) => {
+          return chain()
+            .deleteNode(this.name)
+            .insertContent({
+              type: "paragraph",
+              content: [
+                {
+                  type: "text",
+                  text: content,
+                },
+              ],
+            })
+            .run();
+        },
+      declineAiResponse:
+        ({ prevContent }: { prevContent: string }) =>
+        ({ chain }) => {
+          return chain()
+            .deleteNode(this.name)
+            .insertContent({
+              type: "paragraph",
+              content: [
+                {
+                  type: "text",
+                  text: prevContent,
+                },
+              ],
+            })
+            .run();
+        },
+    };
+  },
+
+  addNodeView: () => {
+    return ReactNodeViewRenderer(AIResponseNode);
+  },
+});