From 9517b01d9a345fac7c5734858802234279195b6f Mon Sep 17 00:00:00 2001
From: Antoine Dewez <44063631+Zewed@users.noreply.github.com>
Date: Tue, 6 Feb 2024 23:34:50 -0800
Subject: [PATCH] feat(frontend): manage current brain (#2165)

# Description

Please include a summary of the changes and the related issue. Please
also include relevant motivation and context.

## Checklist before requesting a review

Please delete options that are not relevant.

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my code
- [ ] I have commented hard-to-understand areas
- [ ] I have ideally added tests that prove my fix is effective or that
my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged

## Screenshots (if appropriate):
---
 .../ActionsBar/components/ChatInput/index.tsx |  2 +-
 frontend/app/chat/[chatId]/page.tsx           | 24 +++++++++++++++++
 frontend/app/search/page.tsx                  |  2 ++
 frontend/app/studio/page.tsx                  |  2 ++
 frontend/app/user/page.tsx                    |  3 +++
 .../components/CurrentBrain/CurrentBrain.tsx  | 26 ++++++++++++-------
 .../lib/components/PageHeader/PageHeader.tsx  |  1 +
 .../lib/components/ui/Icon/Icon.module.scss   |  4 ---
 .../ui/QuivrButton/QuivrButton.module.scss    | 15 +++++++++++
 .../components/ui/QuivrButton/QuivrButton.tsx | 18 ++++++++++++-
 .../lib/components/ui/SearchBar/SearchBar.tsx |  2 +-
 frontend/lib/helpers/iconList.ts              | 15 ++++++-----
 frontend/lib/types/QuivrButton.ts             |  3 +++
 13 files changed, 95 insertions(+), 22 deletions(-)

diff --git a/frontend/app/chat/[chatId]/components/ActionsBar/components/ChatInput/index.tsx b/frontend/app/chat/[chatId]/components/ActionsBar/components/ChatInput/index.tsx
index 95352308078b..ee83f747984a 100644
--- a/frontend/app/chat/[chatId]/components/ActionsBar/components/ChatInput/index.tsx
+++ b/frontend/app/chat/[chatId]/components/ActionsBar/components/ChatInput/index.tsx
@@ -33,7 +33,7 @@ export const ChatInput = (): JSX.Element => {
         }}
       >
         <div className={styles.chat_container}>
-          <CurrentBrain />
+          <CurrentBrain allowingRemoveBrain={false} />
           <div
             className={`
             ${styles.chat_wrapper}
diff --git a/frontend/app/chat/[chatId]/page.tsx b/frontend/app/chat/[chatId]/page.tsx
index 676468c58b3a..089787b4ed6e 100644
--- a/frontend/app/chat/[chatId]/page.tsx
+++ b/frontend/app/chat/[chatId]/page.tsx
@@ -1,9 +1,14 @@
 "use client";
 
+import { UUID } from "crypto";
+import { useEffect } from "react";
+
 import { AddBrainModal } from "@/lib/components/AddBrainModal";
 import { useBrainCreationContext } from "@/lib/components/AddBrainModal/components/AddBrainSteps/brainCreation-provider";
 import PageHeader from "@/lib/components/PageHeader/PageHeader";
 import { UploadDocumentModal } from "@/lib/components/UploadDocumentModal/UploadDocumentModal";
+import { useChatContext } from "@/lib/context";
+import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
 import { useKnowledgeToFeedContext } from "@/lib/context/KnowledgeToFeedProvider/hooks/useKnowledgeToFeedContext";
 import { useDevice } from "@/lib/hooks/useDevice";
 import { useCustomDropzone } from "@/lib/hooks/useDropzone";
@@ -23,6 +28,9 @@ const SelectedChatPage = (): JSX.Element => {
   const { setShouldDisplayFeedCard } = useKnowledgeToFeedContext();
   const { setIsBrainCreationModalOpened } = useBrainCreationContext();
 
+  const { currentBrain, setCurrentBrainId } = useBrainContext();
+  const { messages } = useChatContext();
+
   useChatNotificationsSync();
 
   const buttons: ButtonType[] = [
@@ -32,6 +40,7 @@ const SelectedChatPage = (): JSX.Element => {
       onClick: () => {
         setIsBrainCreationModalOpened(true);
       },
+      iconName: "brain",
     },
     {
       label: "Add knowledge",
@@ -39,9 +48,24 @@ const SelectedChatPage = (): JSX.Element => {
       onClick: () => {
         setShouldDisplayFeedCard(true);
       },
+      iconName: "upload",
+    },
+    {
+      label: "Manage current brain",
+      color: "primary",
+      onClick: () => {
+        window.location.href = `/studio/${currentBrain?.id}`;
+      },
+      iconName: "edit",
     },
   ];
 
+  useEffect(() => {
+    if (!currentBrain && messages.length > 0) {
+      setCurrentBrainId(messages[messages.length - 1].brain_id as UUID);
+    }
+  }, [messages]);
+
   return (
     <div className={styles.main_container}>
       <div className={styles.page_header}>
diff --git a/frontend/app/search/page.tsx b/frontend/app/search/page.tsx
index b33d900b980b..886a352b57ea 100644
--- a/frontend/app/search/page.tsx
+++ b/frontend/app/search/page.tsx
@@ -34,6 +34,7 @@ const Search = (): JSX.Element => {
       onClick: () => {
         setIsBrainCreationModalOpened(true);
       },
+      iconName: "brain",
     },
     {
       label: "Add knowledge",
@@ -41,6 +42,7 @@ const Search = (): JSX.Element => {
       onClick: () => {
         setShouldDisplayFeedCard(true);
       },
+      iconName: "upload",
     },
   ];
 
diff --git a/frontend/app/studio/page.tsx b/frontend/app/studio/page.tsx
index 647c304cbd26..257bce8b99ed 100644
--- a/frontend/app/studio/page.tsx
+++ b/frontend/app/studio/page.tsx
@@ -42,6 +42,7 @@ const Studio = (): JSX.Element => {
       onClick: () => {
         setIsBrainCreationModalOpened(true);
       },
+      iconName: "brain",
     },
     {
       label: "Add knowledge",
@@ -49,6 +50,7 @@ const Studio = (): JSX.Element => {
       onClick: () => {
         setShouldDisplayFeedCard(true);
       },
+      iconName: "upload",
     },
   ];
 
diff --git a/frontend/app/user/page.tsx b/frontend/app/user/page.tsx
index 07364eb115c7..1d4aaea17b6c 100644
--- a/frontend/app/user/page.tsx
+++ b/frontend/app/user/page.tsx
@@ -38,6 +38,7 @@ const UserPage = (): JSX.Element => {
     onClick: () => {
       setIsLogoutModalOpened(true);
     },
+    iconName: "logout",
   };
   const userTabs: Tab[] = [
     {
@@ -91,12 +92,14 @@ const UserPage = (): JSX.Element => {
               onClick={() => setIsLogoutModalOpened(false)}
               color="primary"
               label={t("cancel", { ns: "logout" })}
+              iconName="close"
             ></QuivrButton>
             <QuivrButton
               isLoading={isLoggingOut}
               color="dangerous"
               onClick={() => void handleLogout()}
               label={t("logoutButton")}
+              iconName="logout"
             ></QuivrButton>
           </div>
         </div>
diff --git a/frontend/lib/components/CurrentBrain/CurrentBrain.tsx b/frontend/lib/components/CurrentBrain/CurrentBrain.tsx
index 33f7f43ff18f..bcd5cc1d8d23 100644
--- a/frontend/lib/components/CurrentBrain/CurrentBrain.tsx
+++ b/frontend/lib/components/CurrentBrain/CurrentBrain.tsx
@@ -4,7 +4,13 @@ import styles from "./CurrentBrain.module.scss";
 
 import { Icon } from "../ui/Icon/Icon";
 
-export const CurrentBrain = (): JSX.Element => {
+interface CurrentBrainProps {
+  allowingRemoveBrain: boolean;
+}
+
+export const CurrentBrain = ({
+  allowingRemoveBrain,
+}: CurrentBrainProps): JSX.Element => {
   const { currentBrain, setCurrentBrainId } = useBrainContext();
 
   const removeCurrentBrain = (): void => {
@@ -25,14 +31,16 @@ export const CurrentBrain = (): JSX.Element => {
             <span className={styles.brain_name}>{currentBrain.name}</span>
           </div>
         </div>
-        <div
-          onClick={(event) => {
-            event.nativeEvent.stopImmediatePropagation();
-            removeCurrentBrain();
-          }}
-        >
-          <Icon size="normal" name="close" color="black" handleHover={true} />
-        </div>
+        {allowingRemoveBrain && (
+          <div
+            onClick={(event) => {
+              event.nativeEvent.stopImmediatePropagation();
+              removeCurrentBrain();
+            }}
+          >
+            <Icon size="normal" name="close" color="black" handleHover={true} />
+          </div>
+        )}
       </div>
     </div>
   );
diff --git a/frontend/lib/components/PageHeader/PageHeader.tsx b/frontend/lib/components/PageHeader/PageHeader.tsx
index b48ec7b71c6f..8632c464e840 100644
--- a/frontend/lib/components/PageHeader/PageHeader.tsx
+++ b/frontend/lib/components/PageHeader/PageHeader.tsx
@@ -32,6 +32,7 @@ export const PageHeader = ({
             label={button.label}
             onClick={button.onClick}
             color={button.color}
+            iconName={button.iconName}
           />
         ))}
       </div>
diff --git a/frontend/lib/components/ui/Icon/Icon.module.scss b/frontend/lib/components/ui/Icon/Icon.module.scss
index 543c84bc88ba..570dbc10de14 100644
--- a/frontend/lib/components/ui/Icon/Icon.module.scss
+++ b/frontend/lib/components/ui/Icon/Icon.module.scss
@@ -51,10 +51,6 @@
 
 .white {
   color: Colors.$white;
-
-  &:hover {
-    color: Colors.$accent;
-  }
 }
 
 .dangerous {
diff --git a/frontend/lib/components/ui/QuivrButton/QuivrButton.module.scss b/frontend/lib/components/ui/QuivrButton/QuivrButton.module.scss
index 4da1ce7387f3..4920283f8c4a 100644
--- a/frontend/lib/components/ui/QuivrButton/QuivrButton.module.scss
+++ b/frontend/lib/components/ui/QuivrButton/QuivrButton.module.scss
@@ -1,5 +1,6 @@
 @use "@/styles/Colors.module.scss";
 @use "@/styles/Radius.module.scss";
+@use "@/styles/ScreenSizes.module.scss";
 @use "@/styles/Spacings.module.scss";
 @use "@/styles/Typography.module.scss";
 
@@ -11,6 +12,7 @@
   border-radius: Radius.$normal;
   border: 1.5px solid transparent;
   cursor: pointer;
+  display: flex;
 
   &.primary {
     border-color: Colors.$primary;
@@ -32,3 +34,16 @@
     }
   }
 }
+
+.icon_label {
+  display: flex;
+  flex-direction: row;
+  gap: Spacings.$spacing02;
+  align-items: center;
+
+  @media (max-width: ScreenSizes.$small) {
+    .label {
+      display: none;
+    }
+  }
+}
diff --git a/frontend/lib/components/ui/QuivrButton/QuivrButton.tsx b/frontend/lib/components/ui/QuivrButton/QuivrButton.tsx
index 3d49bf9b8ea7..f71bbdd054a4 100644
--- a/frontend/lib/components/ui/QuivrButton/QuivrButton.tsx
+++ b/frontend/lib/components/ui/QuivrButton/QuivrButton.tsx
@@ -1,7 +1,10 @@
+import { useState } from "react";
+
 import { ButtonType } from "@/lib/types/QuivrButton";
 
 import styles from "./QuivrButton.module.scss";
 
+import { Icon } from "../Icon/Icon";
 import { LoaderIcon } from "../LoaderIcon/LoaderIcon";
 
 export const QuivrButton = ({
@@ -9,14 +12,27 @@ export const QuivrButton = ({
   label,
   color,
   isLoading,
+  iconName,
 }: ButtonType): JSX.Element => {
+  const [hovered, setHovered] = useState<boolean>(false);
+
   return (
     <div
       className={`${styles.button_wrapper} ${styles[color]}`}
       onClick={onClick}
+      onMouseEnter={() => setHovered(true)}
+      onMouseLeave={() => setHovered(false)}
     >
       {!isLoading ? (
-        <span>{label}</span>
+        <div className={styles.icon_label}>
+          <Icon
+            name={iconName}
+            size="normal"
+            color={hovered ? "white" : color}
+            handleHover={false}
+          />
+          <span className={styles.label}>{label}</span>
+        </div>
       ) : (
         <LoaderIcon color="black" size="small" />
       )}
diff --git a/frontend/lib/components/ui/SearchBar/SearchBar.tsx b/frontend/lib/components/ui/SearchBar/SearchBar.tsx
index 957790b9a739..28e63e5fb6a1 100644
--- a/frontend/lib/components/ui/SearchBar/SearchBar.tsx
+++ b/frontend/lib/components/ui/SearchBar/SearchBar.tsx
@@ -56,7 +56,7 @@ export const SearchBar = ({
       ${currentBrain ? styles.with_brain : ""}
       `}
     >
-      <CurrentBrain />
+      <CurrentBrain allowingRemoveBrain={true} />
       <div
         className={`
       ${styles.editor_wrapper}
diff --git a/frontend/lib/helpers/iconList.ts b/frontend/lib/helpers/iconList.ts
index dbff22d630f1..986556e1b8f1 100644
--- a/frontend/lib/helpers/iconList.ts
+++ b/frontend/lib/helpers/iconList.ts
@@ -1,6 +1,6 @@
 import { AiOutlineLoading3Quarters } from "react-icons/ai";
-import { BsArrowRightShort } from "react-icons/bs";
-import { CiChat1, CiFlag1 } from "react-icons/ci";
+import { BsArrowRightShort, BsChatLeftText } from "react-icons/bs";
+import { CiFlag1 } from "react-icons/ci";
 import {
   FaCheck,
   FaCheckCircle,
@@ -9,9 +9,12 @@ import {
   FaRegUserCircle,
   FaUnlock,
 } from "react-icons/fa";
-import { FaArrowUpFromBracket } from "react-icons/fa6";
 import { IoIosAdd, IoMdClose, IoMdLogOut } from "react-icons/io";
-import { IoHomeOutline, IoSettingsSharp } from "react-icons/io5";
+import {
+  IoArrowUpCircleOutline,
+  IoHomeOutline,
+  IoSettingsSharp,
+} from "react-icons/io5";
 import { IconType } from "react-icons/lib";
 import {
   LuBrain,
@@ -38,7 +41,7 @@ export const iconList: { [name: string]: IconType } = {
   addWithoutCircle: IoIosAdd,
   brain: LuBrain,
   brainCircuit: LuBrainCircuit,
-  chat: CiChat1,
+  chat: BsChatLeftText,
   check: FaCheck,
   checkCircle: FaCheckCircle,
   chevronDown: LuChevronDown,
@@ -50,7 +53,7 @@ export const iconList: { [name: string]: IconType } = {
   email: MdAlternateEmail,
   file: LuFile,
   flag: CiFlag1,
-  followUp: FaArrowUpFromBracket,
+  followUp: IoArrowUpCircleOutline,
   graph: VscGraph,
   hastag: RiHashtag,
   history: MdHistory,
diff --git a/frontend/lib/types/QuivrButton.ts b/frontend/lib/types/QuivrButton.ts
index b57f8adbea27..6d56f284e565 100644
--- a/frontend/lib/types/QuivrButton.ts
+++ b/frontend/lib/types/QuivrButton.ts
@@ -1,8 +1,11 @@
 import { Color } from "./Colors";
 
+import { iconList } from "../helpers/iconList";
+
 export interface ButtonType {
   label: string;
   color: Color;
   isLoading?: boolean;
+  iconName: keyof typeof iconList;
   onClick: () => void;
 }