feat: add prompt trigger to mention input

QuivrHQ · Aug 22, 2023 · 1fb39df · 1fb39df
1 parent 0ca25e2
commit 1fb39df
Show file tree

Hide file tree

Showing 12 changed files with 159 additions and 88 deletions.
diff --git a/frontend/app/App.tsx b/frontend/app/App.tsx
@@ -7,17 +7,19 @@ import { NavBar } from "@/lib/components/NavBar";
 import { TrackingWrapper } from "@/lib/components/TrackingWrapper";
 import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
 import { useSupabase } from "@/lib/context/SupabaseProvider";
-import '../lib/config/LocaleConfig/i18n'
 import { UpdateMetadata } from "@/lib/helpers/updateMetadata";
+import "../lib/config/LocaleConfig/i18n";
 
 // This wrapper is used to make effect calls at a high level in app rendering.
 export const App = ({ children }: PropsWithChildren): JSX.Element => {
-  const { fetchAllBrains, fetchAndSetActiveBrain } = useBrainContext();
+  const { fetchAllBrains, fetchAndSetActiveBrain, fetchPublicPrompts } =
+    useBrainContext();
   const { session } = useSupabase();
 
   useEffect(() => {
     void fetchAllBrains();
     void fetchAndSetActiveBrain();
+    void fetchPublicPrompts();
   }, [session?.user]);
 
   return (

diff --git a/...ents/ChatInput/components/ChatBar/components/MentionInput/hooks/helpers/MentionPlugin.tsx b/...ents/ChatInput/components/ChatBar/components/MentionInput/hooks/helpers/MentionPlugin.tsx
@@ -1,9 +1,10 @@
 import createMentionPlugin from "@draft-js-plugins/mention";
 import { useMemo } from "react";
 
+import { MentionTriggerType } from "@/app/chat/[chatId]/components/ActionsBar/types";
 import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
 
-import { BrainMentionItem } from "../../../BrainMentionItem";
+import { MentionItem } from "../../../MentionItem";
 
 interface MentionPluginProps {
   removeMention: (entityKeyToRemove: string) => void;
@@ -12,20 +13,26 @@ interface MentionPluginProps {
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
 export const useMentionPlugin = (props: MentionPluginProps) => {
   const { removeMention } = props;
-  const { setCurrentBrainId } = useBrainContext();
+  const { setCurrentBrainId, setCurrentPromptId } = useBrainContext();
 
   const { MentionSuggestions, plugins } = useMemo(() => {
     const mentionPlugin = createMentionPlugin({
-      mentionComponent: ({ entityKey, mention: { name } }) => (
-        <BrainMentionItem
+      mentionComponent: ({ entityKey, mention: { name, trigger } }) => (
+        <MentionItem
           text={name}
           onRemove={() => {
-            setCurrentBrainId(null);
+            if (trigger === "@") {
+              setCurrentBrainId(null);
+            }
+            if (trigger === "#") {
+              setCurrentPromptId(null);
+            }
             removeMention(entityKey);
           }}
+          trigger={trigger as MentionTriggerType}
         />
       ),
-
+      mentionTrigger: ["@", "#"],
       popperOptions: {
         placement: "top-end",
         modifiers: [

diff --git a/...onents/ChatInput/components/ChatBar/components/MentionInput/hooks/helpers/MentionState.ts b/...onents/ChatInput/components/ChatBar/components/MentionInput/hooks/helpers/MentionState.ts
@@ -1,26 +1,32 @@
+/* eslint-disable max-lines */
+import { MentionData } from "@draft-js-plugins/mention";
 import { EditorState } from "draft-js";
 import { useEffect, useState } from "react";
 
-import { MentionTriggerType } from "@/app/chat/[chatId]/components/ActionsBar/types";
+import {
+  mentionTriggers,
+  MentionTriggerType,
+} from "@/app/chat/[chatId]/components/ActionsBar/types";
 import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
 
 import { MentionInputMentionsType, TriggerMap } from "../../../../types";
 import { mapMinimalBrainToMentionData } from "../../utils/mapMinimalBrainToMentionData";
+import { mapPromptToMentionData } from "../../utils/mapPromptToMentionData";
 
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
 export const useMentionState = () => {
-  const { allBrains } = useBrainContext();
+  const { allBrains, publicPrompts } = useBrainContext();
+
   const [editorState, legacySetEditorState] = useState(() =>
     EditorState.createEmpty()
   );
 
   const [mentionItems, setMentionItems] = useState<MentionInputMentionsType>({
-    "@": allBrains.map((brain) => ({ ...brain, value: brain.name })),
+    "@": allBrains.map(mapMinimalBrainToMentionData),
+    "#": publicPrompts.map(mapPromptToMentionData),
   });
 
-  const [suggestions, setSuggestions] = useState(
-    mapMinimalBrainToMentionData(mentionItems["@"])
-  );
+  const [suggestions, setSuggestions] = useState<MentionData[]>([]);
 
   const setEditorState = (newState: EditorState) => {
     const currentSelection = newState.getSelection();
@@ -35,22 +41,20 @@ export const useMentionState = () => {
   const getEditorCurrentMentions = (): TriggerMap[] => {
     const contentState = editorState.getCurrentContent();
     const plainText = contentState.getPlainText();
-    const mentionTriggers = Object.keys(mentionItems);
+    console.log({ plainText });
 
     const mentionTexts: TriggerMap[] = [];
 
     mentionTriggers.forEach((trigger) => {
-      if (trigger === "@") {
-        mentionItems["@"].forEach((item) => {
-          const mentionText = `${trigger}${item.name}`;
-          if (plainText.includes(mentionText)) {
-            mentionTexts.push({
-              trigger: trigger as MentionTriggerType,
-              content: item.name,
-            });
-          }
-        });
-      }
+      mentionItems[trigger].forEach((item) => {
+        const mentionText = `${trigger}${item.name}`;
+        if (plainText.includes(mentionText)) {
+          mentionTexts.push({
+            trigger: trigger,
+            content: item.name,
+          });
+        }
+      });
     });
 
     return mentionTexts;
@@ -61,8 +65,8 @@ export const useMentionState = () => {
   ): string => {
     const contentState = editorCurrentState.getCurrentContent();
     let plainText = contentState.getPlainText();
-    Object.keys(mentionItems).forEach((trigger) => {
-      if (trigger === "@") {
+    (Object.keys(mentionItems) as MentionTriggerType[]).forEach((trigger) => {
+      if (mentionTriggers.includes(trigger)) {
         mentionItems[trigger].forEach((item) => {
           const regex = new RegExp(`${trigger}${item.name}`, "g");
           plainText = plainText.replace(regex, "");
@@ -76,15 +80,17 @@ export const useMentionState = () => {
   useEffect(() => {
     setMentionItems({
       ...mentionItems,
-      "@": [
-        ...allBrains.map((brain) => ({
-          ...brain,
-          value: brain.name,
-        })),
-      ],
+      "@": allBrains.map(mapMinimalBrainToMentionData),
     });
   }, [allBrains]);
 
+  useEffect(() => {
+    setMentionItems({
+      ...mentionItems,
+      "#": publicPrompts.map(mapPromptToMentionData),
+    });
+  }, [publicPrompts]);
+
   return {
     editorState,
     setEditorState,
@@ -94,5 +100,6 @@ export const useMentionState = () => {
     suggestions,
     getEditorCurrentMentions,
     getEditorTextWithoutMentions,
+    publicPrompts,
   };
 };
diff --git a/...onents/ChatInput/components/ChatBar/components/MentionInput/hooks/helpers/MentionUtils.ts b/...onents/ChatInput/components/ChatBar/components/MentionInput/hooks/helpers/MentionUtils.ts
@@ -1,14 +1,22 @@
 import { addMention, MentionData } from "@draft-js-plugins/mention";
 import { EditorState } from "draft-js";
 
-import { MentionTriggerType } from "@/app/chat/[chatId]/components/ActionsBar/types";
+import {
+  mentionTriggers,
+  MentionTriggerType,
+} from "@/app/chat/[chatId]/components/ActionsBar/types";
 import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
 
 type MentionUtilsProps = {
   editorState: EditorState;
   setEditorState: (editorState: EditorState) => void;
 };
 
+const mentionsTags = [
+  "mention",
+  ...mentionTriggers.map((trigger) => `${trigger}mention`),
+];
+
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
 export const useMentionUtils = (props: MentionUtilsProps) => {
   const { editorState, setEditorState } = props;
@@ -18,7 +26,7 @@ export const useMentionUtils = (props: MentionUtilsProps) => {
     const contentState = editorState.getCurrentContent();
     const entity = contentState.getEntity(entityKeyToRemove);
 
-    if (entity.getType() === "mention") {
+    if (mentionsTags.includes(entity.getType())) {
       const newContentState = contentState.replaceEntityData(
         entityKeyToRemove,
         {}
@@ -37,9 +45,10 @@ export const useMentionUtils = (props: MentionUtilsProps) => {
 
   const insertMention = (
     mention: MentionData,
-    trigger: MentionTriggerType,
     customEditorState?: EditorState
   ): EditorState => {
+    const trigger = mention.trigger as MentionTriggerType;
+
     const editorStateWithMention = addMention(
       customEditorState ?? editorState,
       mention,

diff --git a/...components/ChatInput/components/ChatBar/components/MentionInput/hooks/useMentionInput.tsx b/...components/ChatInput/components/ChatBar/components/MentionInput/hooks/useMentionInput.tsx
@@ -13,10 +13,9 @@ import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainConte
 import { useMentionPlugin } from "./helpers/MentionPlugin";
 import { useMentionState } from "./helpers/MentionState";
 import { useMentionUtils } from "./helpers/MentionUtils";
-import { mapMinimalBrainToMentionData } from "../utils/mapMinimalBrainToMentionData";
-
 import "@draft-js-plugins/mention/lib/plugin.css";
 import "draft-js/dist/Draft.css";
+import { mapMinimalBrainToMentionData } from "../utils/mapMinimalBrainToMentionData";
 
 type UseMentionInputProps = {
   message: string;
@@ -30,7 +29,13 @@ export const useMentionInput = ({
   onSubmit,
   setMessage,
 }: UseMentionInputProps) => {
-  const { allBrains, currentBrainId, setCurrentBrainId } = useBrainContext();
+  const {
+    allBrains,
+    currentBrainId,
+    currentPromptId,
+    setCurrentBrainId,
+    setCurrentPromptId,
+  } = useBrainContext();
 
   const {
     editorState,
@@ -41,6 +46,7 @@ export const useMentionInput = ({
     suggestions,
     getEditorCurrentMentions,
     getEditorTextWithoutMentions,
+    publicPrompts,
   } = useMentionState();
 
   const { removeMention, insertMention } = useMentionUtils({
@@ -64,7 +70,13 @@ export const useMentionInput = ({
   }, []);
 
   const onAddMention = (mention: MentionData) => {
-    setCurrentBrainId(mention.id as UUID);
+    if (mention.trigger === "#") {
+      setCurrentPromptId(mention.id as UUID);
+    }
+
+    if (mention.trigger === "@") {
+      setCurrentBrainId(mention.id as UUID);
+    }
   };
 
   const onSearchChange = ({
@@ -74,18 +86,18 @@ export const useMentionInput = ({
     trigger: string;
     value: string;
   }) => {
-    if (currentBrainId !== null) {
+    if (currentBrainId !== null && trigger === "@") {
       setSuggestions([]);
 
       return;
     }
-    setSuggestions(
-      defaultSuggestionsFilter(
-        value,
-        mapMinimalBrainToMentionData(mentionItems["@"]),
-        trigger
-      )
-    );
+    if (currentPromptId !== null && trigger === "#") {
+      setSuggestions([]);
+
+      return;
+    }
+
+    setSuggestions(defaultSuggestionsFilter(value, mentionItems, trigger));
   };
 
   const insertCurrentBrainAsMention = (): void => {
@@ -94,25 +106,28 @@ export const useMentionInput = ({
     );
 
     if (mention !== undefined) {
-      insertMention(mention, "@");
+      insertMention(mention);
+    }
+  };
+  const insertCurrentPromptAsMention = (): void => {
+    const mention = mentionItems["#"].find(
+      (item) => item.id === currentPromptId
+    );
+
+    if (mention !== undefined) {
+      insertMention(mention);
     }
   };
 
   const resetEditorContent = () => {
     const currentMentions = getEditorCurrentMentions();
     let newEditorState = EditorState.createEmpty();
     currentMentions.forEach((mention) => {
-      if (mention.trigger === "@") {
-        const correspondingMention = mentionItems["@"].find(
-          (item) => item.name === mention.content
-        );
-        if (correspondingMention !== undefined) {
-          newEditorState = insertMention(
-            correspondingMention,
-            mention.trigger,
-            newEditorState
-          );
-        }
+      const correspondingMention = mentionItems[mention.trigger].find(
+        (item) => item.name === mention.content
+      );
+      if (correspondingMention !== undefined) {
+        newEditorState = insertMention(correspondingMention, newEditorState);
       }
     });
     setEditorState(newEditorState);
@@ -144,19 +159,10 @@ export const useMentionInput = ({
     }
   }, [message]);
 
-  useEffect(() => {
-    setSuggestions(mapMinimalBrainToMentionData(mentionItems["@"]));
-  }, [mentionItems]);
-
   useEffect(() => {
     setMentionItems({
       ...mentionItems,
-      "@": [
-        ...allBrains.map((brain) => ({
-          ...brain,
-          value: brain.name,
-        })),
-      ],
+      "@": allBrains.map(mapMinimalBrainToMentionData),
     });
   }, [allBrains]);
 
@@ -178,8 +184,13 @@ export const useMentionInput = ({
     const contentState = editorState.getCurrentContent();
     const plainText = contentState.getPlainText();
 
-    if (plainText === "" && currentBrainId !== null) {
-      insertCurrentBrainAsMention();
+    if (plainText === "") {
+      if (currentBrainId !== null) {
+        insertCurrentBrainAsMention();
+      }
+      if (currentPromptId !== null) {
+        insertCurrentPromptAsMention();
+      }
     }
   }, [editorState]);
 
@@ -196,5 +207,6 @@ export const useMentionInput = ({
     insertCurrentBrainAsMention,
     handleEditorChange,
     keyBindingFn,
+    publicPrompts,
   };
 };