Add global assistants context (#2900)

* add global assistants context * nit * minor cleanup * minor clarity * nit
onyx-dot-app · Oct 24, 2024 · 33eabf1 · 33eabf1
1 parent da979e5
commit 33eabf1
Show file tree

Hide file tree

Showing 13 changed files with 265 additions and 187 deletions.
diff --git a/backend/danswer/db/persona.py b/backend/danswer/db/persona.py
@@ -328,7 +328,6 @@ def update_all_personas_display_priority(
 
     for persona in personas:
         persona.display_priority = display_priority_map[persona.id]
-
     db_session.commit()
 
 

diff --git a/web/src/app/admin/assistants/AssistantEditor.tsx b/web/src/app/admin/assistants/AssistantEditor.tsx
@@ -61,6 +61,7 @@ import {
 import { AdvancedOptionsToggle } from "@/components/AdvancedOptionsToggle";
 import { buildImgUrl } from "@/app/chat/files/images/utils";
 import { LlmList } from "@/components/llm/LLMList";
+import { useAssistants } from "@/components/context/AssistantsContext";
 
 function findSearchTool(tools: ToolSnapshot[]) {
   return tools.find((tool) => tool.in_code_tool_id === "SearchTool");
@@ -105,6 +106,7 @@ export function AssistantEditor({
   shouldAddAssistantToUserPreferences?: boolean;
   admin?: boolean;
 }) {
+  const { refreshAssistants } = useAssistants();
   const router = useRouter();
 
   const { popup, setPopup } = usePopup();
@@ -433,6 +435,7 @@ export function AssistantEditor({
                 });
               }
             }
+            await refreshAssistants();
             router.push(
               redirectType === SuccessfulPersonaUpdateRedirectType.ADMIN
                 ? `/admin/assistants?u=${Date.now()}`

diff --git a/web/src/app/admin/assistants/PersonaTable.tsx b/web/src/app/admin/assistants/PersonaTable.tsx
@@ -5,7 +5,7 @@ import { Persona } from "./interfaces";
 import { useRouter } from "next/navigation";
 import { CustomCheckbox } from "@/components/CustomCheckbox";
 import { usePopup } from "@/components/admin/connectors/Popup";
-import { useState, useMemo } from "react";
+import { useState, useMemo, useEffect } from "react";
 import { UniqueIdentifier } from "@dnd-kit/core";
 import { DraggableTable } from "@/components/table/DraggableTable";
 import {
@@ -16,6 +16,7 @@ import {
 import { FiEdit2 } from "react-icons/fi";
 import { TrashIcon } from "@/components/icons/icons";
 import { useUser } from "@/components/user/UserProvider";
+import { useAssistants } from "@/components/context/AssistantsContext";
 
 function PersonaTypeDisplay({ persona }: { persona: Persona }) {
   if (persona.builtin_persona) {
@@ -37,43 +38,36 @@ function PersonaTypeDisplay({ persona }: { persona: Persona }) {
   return <Text>Personal {persona.owner && <>({persona.owner.email})</>}</Text>;
 }
 
-export function PersonasTable({
-  allPersonas,
-  editablePersonas,
-}: {
-  allPersonas: Persona[];
-  editablePersonas: Persona[];
-}) {
+export function PersonasTable() {
   const router = useRouter();
   const { popup, setPopup } = usePopup();
-
-  const { isLoadingUser, isAdmin } = useUser();
+  const { refreshUser, isLoadingUser, isAdmin } = useUser();
+  const {
+    allAssistants: assistants,
+    refreshAssistants,
+    editablePersonas,
+  } = useAssistants();
 
   const editablePersonaIds = useMemo(() => {
     return new Set(editablePersonas.map((p) => p.id.toString()));
   }, [editablePersonas]);
 
-  const sortedPersonas = useMemo(() => {
+  const [finalPersonas, setFinalPersonas] = useState<Persona[]>([]);
+
+  useEffect(() => {
     const editable = editablePersonas.sort(personaComparator);
-    const nonEditable = allPersonas
+    const nonEditable = assistants
       .filter((p) => !editablePersonaIds.has(p.id.toString()))
       .sort(personaComparator);
-    return [...editable, ...nonEditable];
-  }, [allPersonas, editablePersonas]);
-
-  const [finalPersonas, setFinalPersonas] = useState<string[]>(
-    sortedPersonas.map((persona) => persona.id.toString())
-  );
-  const finalPersonaValues = finalPersonas
-    .filter((id) => new Set(allPersonas.map((p) => p.id.toString())).has(id))
-    .map((id) => {
-      return sortedPersonas.find(
-        (persona) => persona.id.toString() === id
-      ) as Persona;
-    });
+    setFinalPersonas([...editable, ...nonEditable]);
+  }, [editablePersonas, assistants, editablePersonaIds]);
 
   const updatePersonaOrder = async (orderedPersonaIds: UniqueIdentifier[]) => {
-    setFinalPersonas(orderedPersonaIds.map((id) => id.toString()));
+    const reorderedAssistants = orderedPersonaIds.map(
+      (id) => assistants.find((assistant) => assistant.id.toString() === id)!
+    );
+
+    setFinalPersonas(reorderedAssistants);
 
     const displayPriorityMap = new Map<UniqueIdentifier, number>();
     orderedPersonaIds.forEach((personaId, ind) => {
@@ -89,13 +83,19 @@ export function PersonasTable({
         display_priority_map: Object.fromEntries(displayPriorityMap),
       }),
     });
+
     if (!response.ok) {
       setPopup({
         type: "error",
         message: `Failed to update persona order - ${await response.text()}`,
       });
+      setFinalPersonas(assistants);
       router.refresh();
+      return;
     }
+
+    await refreshAssistants();
+    await refreshUser();
   };
 
   if (isLoadingUser) {
@@ -115,8 +115,8 @@ export function PersonasTable({
       <DraggableTable
         headers={["Name", "Description", "Type", "Is Visible", "Delete"]}
         isAdmin={isAdmin}
-        rows={finalPersonaValues.map((persona) => {
-          const isEditable = editablePersonaIds.has(persona.id.toString());
+        rows={finalPersonas.map((persona) => {
+          const isEditable = editablePersonas.includes(persona);
           return {
             id: persona.id.toString(),
             cells: [

diff --git a/web/src/app/admin/assistants/page.tsx b/web/src/app/admin/assistants/page.tsx
@@ -2,33 +2,10 @@ import { PersonasTable } from "./PersonaTable";
 import { FiPlusSquare } from "react-icons/fi";
 import Link from "next/link";
 import { Divider, Text, Title } from "@tremor/react";
-import { fetchSS } from "@/lib/utilsSS";
-import { ErrorCallout } from "@/components/ErrorCallout";
-import { Persona } from "./interfaces";
 import { AssistantsIcon } from "@/components/icons/icons";
 import { AdminPageTitle } from "@/components/admin/Title";
 
 export default async function Page() {
-  const allPersonaResponse = await fetchSS("/admin/persona");
-  const editablePersonaResponse = await fetchSS(
-    "/admin/persona?get_editable=true"
-  );
-
-  if (!allPersonaResponse.ok || !editablePersonaResponse.ok) {
-    return (
-      <ErrorCallout
-        errorTitle="Something went wrong :("
-        errorMsg={`Failed to fetch personas - ${
-          (await allPersonaResponse.text()) ||
-          (await editablePersonaResponse.text())
-        }`}
-      />
-    );
-  }
-
-  const allPersonas = (await allPersonaResponse.json()) as Persona[];
-  const editablePersonas = (await editablePersonaResponse.json()) as Persona[];
-
   return (
     <div className="mx-auto container">
       <AdminPageTitle icon={<AssistantsIcon size={32} />} title="Assistants" />
@@ -64,10 +41,7 @@ export default async function Page() {
         <Divider />
 
         <Title>Existing Assistants</Title>
-        <PersonasTable
-          allPersonas={allPersonas}
-          editablePersonas={editablePersonas}
-        />
+        <PersonasTable />
       </div>
     </div>
   );

diff --git a/web/src/app/assistants/gallery/page.tsx b/web/src/app/assistants/gallery/page.tsx
@@ -22,33 +22,24 @@ export default async function GalleryPage({
   const {
     user,
     chatSessions,
-    assistants,
     folders,
     openedFolders,
     shouldShowWelcomeModal,
     toggleSidebar,
-    hasAnyConnectors,
-    hasImageCompatibleModel,
   } = data;
 
   return (
     <>
-      <AssistantsProvider
-        initialAssistants={assistants}
-        hasAnyConnectors={hasAnyConnectors}
-        hasImageCompatibleModel={hasImageCompatibleModel}
-      >
-        {shouldShowWelcomeModal && <WelcomeModal user={user} />}
+      {shouldShowWelcomeModal && <WelcomeModal user={user} />}
 
-        <InstantSSRAutoRefresh />
+      <InstantSSRAutoRefresh />
 
-        <WrappedAssistantsGallery
-          initiallyToggled={toggleSidebar}
-          chatSessions={chatSessions}
-          folders={folders}
-          openedFolders={openedFolders}
-        />
-      </AssistantsProvider>
+      <WrappedAssistantsGallery
+        initiallyToggled={toggleSidebar}
+        chatSessions={chatSessions}
+        folders={folders}
+        openedFolders={openedFolders}
+      />
     </>
   );
 }
diff --git a/web/src/app/assistants/mine/page.tsx b/web/src/app/assistants/mine/page.tsx
@@ -23,20 +23,13 @@ export default async function GalleryPage({
     user,
     chatSessions,
     folders,
-    assistants,
     openedFolders,
     shouldShowWelcomeModal,
     toggleSidebar,
-    hasAnyConnectors,
-    hasImageCompatibleModel,
   } = data;
 
   return (
-    <AssistantsProvider
-      initialAssistants={assistants}
-      hasAnyConnectors={hasAnyConnectors}
-      hasImageCompatibleModel={hasImageCompatibleModel}
-    >
+    <>
       {shouldShowWelcomeModal && <WelcomeModal user={user} />}
 
       <InstantSSRAutoRefresh />
@@ -46,6 +39,6 @@ export default async function GalleryPage({
         folders={folders}
         openedFolders={openedFolders}
       />
-    </AssistantsProvider>
+    </>
   );
 }
diff --git a/web/src/app/chat/page.tsx b/web/src/app/chat/page.tsx
@@ -32,38 +32,29 @@ export default async function Page({
     openedFolders,
     defaultAssistantId,
     shouldShowWelcomeModal,
-    assistants,
     userInputPrompts,
-    hasAnyConnectors,
-    hasImageCompatibleModel,
   } = data;
 
   return (
     <>
       <InstantSSRAutoRefresh />
       {shouldShowWelcomeModal && <WelcomeModal user={user} />}
-      <AssistantsProvider
-        initialAssistants={assistants}
-        hasAnyConnectors={hasAnyConnectors}
-        hasImageCompatibleModel={hasImageCompatibleModel}
+      <ChatProvider
+        value={{
+          chatSessions,
+          availableSources,
+          availableDocumentSets: documentSets,
+          availableTags: tags,
+          llmProviders,
+          folders,
+          openedFolders,
+          userInputPrompts,
+          shouldShowWelcomeModal,
+          defaultAssistantId,
+        }}
       >
-        <ChatProvider
-          value={{
-            chatSessions,
-            availableSources,
-            availableDocumentSets: documentSets,
-            availableTags: tags,
-            llmProviders,
-            folders,
-            openedFolders,
-            userInputPrompts,
-            shouldShowWelcomeModal,
-            defaultAssistantId,
-          }}
-        >
-          <WrappedChat initiallyToggled={toggleSidebar} />
-        </ChatProvider>
-      </AssistantsProvider>
+        <WrappedChat initiallyToggled={toggleSidebar} />
+      </ChatProvider>
     </>
   );
 }
diff --git a/web/src/app/layout.tsx b/web/src/app/layout.tsx
@@ -19,6 +19,8 @@ import { HeaderTitle } from "@/components/header/HeaderTitle";
 import { Logo } from "@/components/Logo";
 import { UserProvider } from "@/components/user/UserProvider";
 import { ProviderContextProvider } from "@/components/chat_search/ProviderContext";
+import { fetchAssistantData } from "@/lib/chat/fetchAssistantdata";
+import { AppProvider } from "@/components/context/AppProvider";
 
 const inter = Inter({
   subsets: ["latin"],
@@ -55,6 +57,10 @@ export default async function RootLayout({
 }) {
   const combinedSettings = await fetchSettingsSS();
 
+  const data = await fetchAssistantData();
+
+  const { assistants, hasAnyConnectors, hasImageCompatibleModel } = data;
+
   const productGating =
     combinedSettings?.settings.product_gating ?? GatingType.NONE;
 
@@ -174,13 +180,14 @@ export default async function RootLayout({
             process.env.THEME_IS_DARK?.toLowerCase() === "true" ? "dark" : ""
           }`}
         >
-          <UserProvider>
-            <ProviderContextProvider>
-              <SettingsProvider settings={combinedSettings}>
-                {children}
-              </SettingsProvider>
-            </ProviderContextProvider>
-          </UserProvider>
+          <AppProvider
+            settings={combinedSettings}
+            assistants={assistants}
+            hasAnyConnectors={hasAnyConnectors}
+            hasImageCompatibleModel={hasImageCompatibleModel}
+          >
+            {children}
+          </AppProvider>
         </div>
       </body>
     </html>
Original file line number	Diff line number	Diff line change
Expand Up		@@ -328,7 +328,6 @@ def update_all_personas_display_priority(

		for persona in personas:
		persona.display_priority = display_priority_map[persona.id]

		db_session.commit()


Expand Down