From a30de693cb6949c97405e550c1d1acd9d93c4da2 Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Wed, 2 Oct 2024 09:15:54 -0700
Subject: [PATCH] Clean, memoized assistant ordering (#2655)

* updated refresh

* memoization and so on

* nit

* build issue
---
 web/src/app/chat/ChatPage.tsx                 | 21 +++++++++-----
 web/src/app/chat/input/ChatInputBar.tsx       |  3 ++
 .../modal/configuration/AssistantsTab.tsx     | 29 +++++++++----------
 web/src/components/user/UserProvider.tsx      |  1 -
 4 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx
index e9676d282..8f9b36b62 100644
--- a/web/src/app/chat/ChatPage.tsx
+++ b/web/src/app/chat/ChatPage.tsx
@@ -50,6 +50,7 @@ import {
   useContext,
   useEffect,
   useLayoutEffect,
+  useMemo,
   useRef,
   useState,
 } from "react";
@@ -157,14 +158,17 @@ export function ChatPage({
   // Useful for determining which session has been loaded (i.e. still on `new, empty session` or `previous session`)
   const loadedIdSessionRef = useRef<number | null>(existingChatSessionId);
 
-  // Assistants
-  const { visibleAssistants, hiddenAssistants: _ } = classifyAssistants(
-    user,
-    availableAssistants
-  );
-  const finalAssistants = user
-    ? orderAssistantsForUser(visibleAssistants, user)
-    : visibleAssistants;
+  // Assistants in order
+  const { finalAssistants } = useMemo(() => {
+    const { visibleAssistants, hiddenAssistants: _ } = classifyAssistants(
+      user,
+      availableAssistants
+    );
+    const finalAssistants = user
+      ? orderAssistantsForUser(visibleAssistants, user)
+      : visibleAssistants;
+    return { finalAssistants };
+  }, [user, availableAssistants]);
 
   const existingChatSessionAssistantId = selectedChatSession?.persona_id;
   const [selectedAssistant, setSelectedAssistant] = useState<
@@ -2410,6 +2414,7 @@ export function ChatPage({
                               handleFileUpload={handleImageUpload}
                               textAreaRef={textAreaRef}
                               chatSessionId={chatSessionIdRef.current!}
+                              refreshUser={refreshUser}
                             />
 
                             {enterpriseSettings &&
diff --git a/web/src/app/chat/input/ChatInputBar.tsx b/web/src/app/chat/input/ChatInputBar.tsx
index 64535d82b..cef0bc49f 100644
--- a/web/src/app/chat/input/ChatInputBar.tsx
+++ b/web/src/app/chat/input/ChatInputBar.tsx
@@ -63,6 +63,7 @@ export function ChatInputBar({
   alternativeAssistant,
   chatSessionId,
   inputPrompts,
+  refreshUser,
 }: {
   showConfigureAPIKey: () => void;
   openModelSettings: () => void;
@@ -86,6 +87,7 @@ export function ChatInputBar({
   handleFileUpload: (files: File[]) => void;
   textAreaRef: React.RefObject<HTMLTextAreaElement>;
   chatSessionId?: number;
+  refreshUser: () => void;
 }) {
   useEffect(() => {
     const textarea = textAreaRef.current;
@@ -532,6 +534,7 @@ export function ChatInputBar({
                       setSelectedAssistant(assistant);
                       close();
                     }}
+                    refreshUser={refreshUser}
                   />
                 )}
                 flexPriority="shrink"
diff --git a/web/src/app/chat/modal/configuration/AssistantsTab.tsx b/web/src/app/chat/modal/configuration/AssistantsTab.tsx
index dcf31138c..1db03cd60 100644
--- a/web/src/app/chat/modal/configuration/AssistantsTab.tsx
+++ b/web/src/app/chat/modal/configuration/AssistantsTab.tsx
@@ -13,25 +13,26 @@ import {
   sortableKeyboardCoordinates,
   verticalListSortingStrategy,
 } from "@dnd-kit/sortable";
-import { CSS } from "@dnd-kit/utilities";
 import { Persona } from "@/app/admin/assistants/interfaces";
 import { LLMProviderDescriptor } from "@/app/admin/configuration/llm/interfaces";
 import { getFinalLLM } from "@/lib/llm/utils";
 import React, { useState } from "react";
 import { updateUserAssistantList } from "@/lib/assistants/updateAssistantPreferences";
 import { DraggableAssistantCard } from "@/components/assistants/AssistantCards";
-import { orderAssistantsForUser } from "@/lib/assistants/utils";
+import { useRouter } from "next/navigation";
 
 export function AssistantsTab({
   selectedAssistant,
   availableAssistants,
   llmProviders,
   onSelect,
+  refreshUser,
 }: {
   selectedAssistant: Persona;
   availableAssistants: Persona[];
   llmProviders: LLMProviderDescriptor[];
   onSelect: (assistant: Persona) => void;
+  refreshUser: () => void;
 }) {
   const [_, llmName] = getFinalLLM(llmProviders, null, null);
   const [assistants, setAssistants] = useState(availableAssistants);
@@ -43,23 +44,21 @@ export function AssistantsTab({
     })
   );
 
-  function handleDragEnd(event: DragEndEvent) {
+  async function handleDragEnd(event: DragEndEvent) {
     const { active, over } = event;
 
     if (over && active.id !== over.id) {
-      setAssistants((items) => {
-        const oldIndex = items.findIndex(
-          (item) => item.id.toString() === active.id
-        );
-        const newIndex = items.findIndex(
-          (item) => item.id.toString() === over.id
-        );
-        const updatedAssistants = arrayMove(items, oldIndex, newIndex);
+      const oldIndex = assistants.findIndex(
+        (item) => item.id.toString() === active.id
+      );
+      const newIndex = assistants.findIndex(
+        (item) => item.id.toString() === over.id
+      );
+      const updatedAssistants = arrayMove(assistants, oldIndex, newIndex);
 
-        updateUserAssistantList(updatedAssistants.map((a) => a.id));
-
-        return updatedAssistants;
-      });
+      setAssistants(updatedAssistants);
+      await updateUserAssistantList(updatedAssistants.map((a) => a.id));
+      refreshUser();
     }
   }
 
diff --git a/web/src/components/user/UserProvider.tsx b/web/src/components/user/UserProvider.tsx
index 67777277c..5e7bb7535 100644
--- a/web/src/components/user/UserProvider.tsx
+++ b/web/src/components/user/UserProvider.tsx
@@ -40,7 +40,6 @@ export function UserProvider({ children }: { children: React.ReactNode }) {
   }, []);
 
   const refreshUser = async () => {
-    setIsLoadingUser(true);
     await fetchUser();
   };