Add global assistants context (#2900)

* add global assistants context

* nit

* minor cleanup

* minor clarity

* nit

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()
 
 

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()}`

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({
+export function PersonasTable() {
-  allPersonas,
-  editablePersonas,
-}: {
-  allPersonas: Persona[];
-  editablePersonas: Persona[];
-}) {
   const router = useRouter();
   const { popup, setPopup } = usePopup();
-
+  const { refreshUser, isLoadingUser, isAdmin } = useUser();
-  const { 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];
+    setFinalPersonas([...editable, ...nonEditable]);
-  }, [allPersonas, editablePersonas]);
+  }, [editablePersonas, assistants, editablePersonaIds]);
-
-  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;
-    });
 
   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) => {
+        rows={finalPersonas.map((persona) => {
-          const isEditable = editablePersonaIds.has(persona.id.toString());
+          const isEditable = editablePersonas.includes(persona);
           return {
             id: persona.id.toString(),
             cells: [

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
+        <PersonasTable />
-          allPersonas={allPersonas}
-          editablePersonas={editablePersonas}
-        />
       </div>
     </div>
   );

web/src/app/assistants/gallery/page.tsx

@@ -22,22 +22,14 @@ 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} />}
 
       <InstantSSRAutoRefresh />
@@ -48,7 +40,6 @@ export default async function GalleryPage({
         folders={folders}
         openedFolders={openedFolders}
       />
-      </AssistantsProvider>
     </>
   );
 }

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>
+    </>
   );
 }

web/src/app/chat/page.tsx

@@ -32,21 +32,13 @@ 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,
@@ -63,7 +55,6 @@ export default async function Page({
       >
         <WrappedChat initiallyToggled={toggleSidebar} />
       </ChatProvider>
-      </AssistantsProvider>
     </>
   );
 }

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>
+          <AppProvider
-            <ProviderContextProvider>
+            settings={combinedSettings}
-              <SettingsProvider settings={combinedSettings}>
+            assistants={assistants}
+            hasAnyConnectors={hasAnyConnectors}
+            hasImageCompatibleModel={hasImageCompatibleModel}
+          >
             {children}
-              </SettingsProvider>
+          </AppProvider>
-            </ProviderContextProvider>
-          </UserProvider>
         </div>
       </body>
     </html>

web/src/app/search/page.tsx

@@ -201,11 +201,6 @@ export default async function Home({
       {/* ChatPopup is a custom popup that displays a admin-specified message on initial user visit. 
       Only used in the EE version of the app. */}
       <ChatPopup />
-      <AssistantsProvider
-        initialAssistants={assistants}
-        hasAnyConnectors={hasAnyConnectors}
-        hasImageCompatibleModel={false}
-      >
       <SearchProvider
         value={{
           querySessions,
@@ -225,8 +220,6 @@ export default async function Home({
           searchTypeDefault={searchTypeDefault}
         />
       </SearchProvider>
-      </AssistantsProvider>
-      s
     </>
   );
 }

web/src/components/context/AppProvider.tsx

@@ -0,0 +1,39 @@
+"use server";
+import { CombinedSettings } from "@/app/admin/settings/interfaces";
+import { UserProvider } from "../user/UserProvider";
+import { ProviderContextProvider } from "../chat_search/ProviderContext";
+import { SettingsProvider } from "../settings/SettingsProvider";
+import { AssistantsProvider } from "./AssistantsContext";
+import { Persona } from "@/app/admin/assistants/interfaces";
+
+interface AppProviderProps {
+  children: React.ReactNode;
+  settings: CombinedSettings;
+  assistants: Persona[];
+  hasAnyConnectors: boolean;
+  hasImageCompatibleModel: boolean;
+}
+
+export const AppProvider = ({
+  children,
+  settings,
+  assistants,
+  hasAnyConnectors,
+  hasImageCompatibleModel,
+}: AppProviderProps) => {
+  return (
+    <UserProvider>
+      <ProviderContextProvider>
+        <SettingsProvider settings={settings}>
+          <AssistantsProvider
+            initialAssistants={assistants}
+            hasAnyConnectors={hasAnyConnectors}
+            hasImageCompatibleModel={hasImageCompatibleModel}
+          >
+            {children}
+          </AssistantsProvider>
+        </SettingsProvider>
+      </ProviderContextProvider>
+    </UserProvider>
+  );
+};

web/src/components/context/AssistantsContext.tsx

@@ -1,5 +1,11 @@
 "use client";
-import React, { createContext, useState, useContext, useMemo } from "react";
+import React, {
+  createContext,
+  useState,
+  useContext,
+  useMemo,
+  useEffect,
+} from "react";
 import { Persona } from "@/app/admin/assistants/interfaces";
 import {
   classifyAssistants,
@@ -15,6 +21,10 @@ interface AssistantsContextProps {
   finalAssistants: Persona[];
   ownedButHiddenAssistants: Persona[];
   refreshAssistants: () => Promise<void>;
+
+  // Admin only
+  editablePersonas: Persona[];
+  allAssistants: Persona[];
 }
 
 const AssistantsContext = createContext<AssistantsContextProps | undefined>(
@@ -35,7 +45,54 @@ export const AssistantsProvider: React.FC<{
   const [assistants, setAssistants] = useState<Persona[]>(
     initialAssistants || []
   );
-  const { user } = useUser();
+  const { user, isLoadingUser, isAdmin } = useUser();
+  const [editablePersonas, setEditablePersonas] = useState<Persona[]>([]);
+
+  useEffect(() => {
+    const fetchEditablePersonas = async () => {
+      if (!isAdmin) {
+        return;
+      }
+
+      try {
+        const response = await fetch("/api/admin/persona?get_editable=true");
+        if (!response.ok) {
+          console.error("Failed to fetch editable personas");
+          return;
+        }
+        const personas = await response.json();
+        setEditablePersonas(personas);
+      } catch (error) {
+        console.error("Error fetching editable personas:", error);
+      }
+    };
+
+    fetchEditablePersonas();
+  }, [isAdmin]);
+
+  const [allAssistants, setAllAssistants] = useState<Persona[]>([]);
+
+  useEffect(() => {
+    const fetchAllAssistants = async () => {
+      if (!isAdmin) {
+        return;
+      }
+
+      try {
+        const response = await fetch("/api/admin/persona");
+        if (!response.ok) {
+          console.error("Failed to fetch all personas");
+          return;
+        }
+        const personas = await response.json();
+        setAllAssistants(personas);
+      } catch (error) {
+        console.error("Error fetching all personas:", error);
+      }
+    };
+
+    fetchAllAssistants();
+  }, [isAdmin]);
 
   const refreshAssistants = async () => {
     try {
@@ -92,7 +149,7 @@ export const AssistantsProvider: React.FC<{
       finalAssistants,
       ownedButHiddenAssistants,
     };
-  }, [user, assistants]);
+  }, [user, assistants, isLoadingUser]);
 
   return (
     <AssistantsContext.Provider
@@ -103,6 +160,8 @@ export const AssistantsProvider: React.FC<{
         finalAssistants,
         ownedButHiddenAssistants,
         refreshAssistants,
+        editablePersonas,
+        allAssistants,
       }}
     >
       {children}

web/src/lib/chat/fetchAssistantdata.ts

@@ -0,0 +1,69 @@
+import { fetchSS } from "@/lib/utilsSS";
+import { CCPairBasicInfo } from "@/lib/types";
+import { Persona } from "@/app/admin/assistants/interfaces";
+import { fetchLLMProvidersSS } from "@/lib/llm/fetchLLMs";
+import { personaComparator } from "@/app/admin/assistants/lib";
+import { fetchAssistantsSS } from "../assistants/fetchAssistantsSS";
+import { checkLLMSupportsImageInput } from "../llm/utils";
+
+interface AssistantData {
+  assistants: Persona[];
+  hasAnyConnectors: boolean;
+  hasImageCompatibleModel: boolean;
+}
+export async function fetchAssistantData(): Promise<AssistantData> {
+  const [assistants, assistantsFetchError] = await fetchAssistantsSS();
+  const ccPairsResponse = await fetchSS("/manage/indexing-status");
+
+  let ccPairs: CCPairBasicInfo[] = [];
+  if (ccPairsResponse?.ok) {
+    ccPairs = await ccPairsResponse.json();
+  } else {
+    console.log(`Failed to fetch connectors - ${ccPairsResponse?.status}`);
+  }
+
+  const hasAnyConnectors = ccPairs.length > 0;
+
+  // if no connectors are setup, only show personas that are pure
+  // passthrough and don't do any retrieval
+  let filteredAssistants = assistants;
+  if (assistantsFetchError) {
+    console.log(`Failed to fetch assistants - ${assistantsFetchError}`);
+  }
+
+  // remove those marked as hidden by an admin
+  filteredAssistants = filteredAssistants.filter(
+    (assistant) => assistant.is_visible
+  );
+
+  if (!hasAnyConnectors) {
+    filteredAssistants = filteredAssistants.filter(
+      (assistant) => assistant.num_chunks === 0
+    );
+  }
+
+  // sort them in priority order
+  filteredAssistants.sort(personaComparator);
+
+  const llmProviders = await fetchLLMProvidersSS();
+  const hasImageCompatibleModel = llmProviders.some(
+    (provider) =>
+      provider.provider === "openai" ||
+      provider.model_names.some((model) => checkLLMSupportsImageInput(model))
+  );
+
+  if (!hasImageCompatibleModel) {
+    filteredAssistants = filteredAssistants.filter(
+      (assistant) =>
+        !assistant.tools.some(
+          (tool) => tool.in_code_tool_id === "ImageGenerationTool"
+        )
+    );
+  }
+
+  return {
+    assistants: filteredAssistants,
+    hasAnyConnectors,
+    hasImageCompatibleModel,
+  };
+}

web/src/lib/chat/fetchChatData.ts

@@ -37,7 +37,6 @@ interface FetchChatDataResult {
   ccPairs: CCPairBasicInfo[];
   availableSources: ValidSources[];
   documentSets: DocumentSet[];
-  assistants: Persona[];
   tags: Tag[];
   llmProviders: LLMProviderDescriptor[];
   folders: Folder[];
@@ -47,8 +46,6 @@ interface FetchChatDataResult {
   finalDocumentSidebarInitialWidth?: number;
   shouldShowWelcomeModal: boolean;
   userInputPrompts: InputPrompt[];
-  hasAnyConnectors: boolean;
-  hasImageCompatibleModel: boolean;
 }
 
 export async function fetchChatData(searchParams: {
@@ -59,7 +56,6 @@ export async function fetchChatData(searchParams: {
     getCurrentUserSS(),
     fetchSS("/manage/indexing-status"),
     fetchSS("/manage/document-set"),
-    fetchAssistantsSS(),
     fetchSS("/chat/get-user-chat-sessions"),
     fetchSS("/query/valid-tags"),
     fetchLLMProvidersSS(),
@@ -76,7 +72,7 @@ export async function fetchChatData(searchParams: {
     | LLMProviderDescriptor[]
     | [Persona[], string | null]
     | null
-  )[] = [null, null, null, null, null, null, null, null, null, null];
+  )[] = [null, null, null, null, null, null, null, null, null];
   try {
     results = await Promise.all(tasks);
   } catch (e) {
@@ -87,17 +83,13 @@ export async function fetchChatData(searchParams: {
   const user = results[1] as User | null;
   const ccPairsResponse = results[2] as Response | null;
   const documentSetsResponse = results[3] as Response | null;
-  const [rawAssistantsList, assistantsFetchError] = results[4] as [
-    Persona[],
-    string | null,
-  ];
 
-  const chatSessionsResponse = results[5] as Response | null;
+  const chatSessionsResponse = results[4] as Response | null;
 
-  const tagsResponse = results[6] as Response | null;
+  const tagsResponse = results[5] as Response | null;
-  const llmProviders = (results[7] || []) as LLMProviderDescriptor[];
+  const llmProviders = (results[6] || []) as LLMProviderDescriptor[];
-  const foldersResponse = results[8] as Response | null;
+  const foldersResponse = results[7] as Response | null;
-  const userInputPromptsResponse = results[9] as Response | null;
+  const userInputPromptsResponse = results[8] as Response | null;
 
   const authDisabled = authTypeMetadata?.authType === "disabled";
   if (!authDisabled && !user) {
@@ -161,17 +153,6 @@ export async function fetchChatData(searchParams: {
     );
   }
 
-  let assistants = rawAssistantsList;
-  if (assistantsFetchError) {
-    console.log(`Failed to fetch assistants - ${assistantsFetchError}`);
-  }
-  // remove those marked as hidden by an admin
-
-  assistants = assistants.filter((assistant) => assistant.is_visible);
-
-  // sort them in priority order
-  assistants.sort(personaComparator);
-
   let tags: Tag[] = [];
   if (tagsResponse?.ok) {
     tags = (await tagsResponse.json()).tags;
@@ -206,24 +187,6 @@ export async function fetchChatData(searchParams: {
 
   // if no connectors are setup, only show personas that are pure
   // passthrough and don't do any retrieval
-  if (!hasAnyConnectors) {
-    assistants = assistants.filter((assistant) => assistant.num_chunks === 0);
-  }
-
-  const hasImageCompatibleModel = llmProviders.some(
-    (provider) =>
-      provider.provider === "openai" ||
-      provider.model_names.some((model) => checkLLMSupportsImageInput(model))
-  );
-
-  if (!hasImageCompatibleModel) {
-    assistants = assistants.filter(
-      (assistant) =>
-        !assistant.tools.some(
-          (tool) => tool.in_code_tool_id === "ImageGenerationTool"
-        )
-    );
-  }
 
   let folders: Folder[] = [];
   if (foldersResponse?.ok) {
@@ -243,7 +206,6 @@ export async function fetchChatData(searchParams: {
     ccPairs,
     availableSources,
     documentSets,
-    assistants,
     tags,
     llmProviders,
     folders,
@@ -253,7 +215,5 @@ export async function fetchChatData(searchParams: {
     toggleSidebar,
     shouldShowWelcomeModal,
     userInputPrompts,
-    hasAnyConnectors,
-    hasImageCompatibleModel,
   };
 }