Add global assistants context (#2900)

* add global assistants context

* nit

* minor cleanup

* minor clarity

* nit
This commit is contained in:
pablodanswer 2024-10-24 14:27:55 -07:00 committed by GitHub
parent da979e5745
commit 33eabf1b25
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 265 additions and 187 deletions

View File

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

View File

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

View File

@ -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: [

View File

@ -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>
);

View File

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

View File

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

View File

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

View File

@ -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>

View File

@ -201,32 +201,25 @@ 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,
ccPairs,
documentSets,
assistants,
tags,
agenticSearchEnabled,
disabledAgentic: DISABLE_LLM_DOC_RELEVANCE,
initiallyToggled: toggleSidebar,
shouldShowWelcomeModal,
shouldDisplayNoSources: shouldDisplayNoSourcesModal,
}}
>
<SearchProvider
value={{
querySessions,
ccPairs,
documentSets,
assistants,
tags,
agenticSearchEnabled,
disabledAgentic: DISABLE_LLM_DOC_RELEVANCE,
initiallyToggled: toggleSidebar,
shouldShowWelcomeModal,
shouldDisplayNoSources: shouldDisplayNoSourcesModal,
}}
>
<WrappedSearch
initiallyToggled={toggleSidebar}
searchTypeDefault={searchTypeDefault}
/>
</SearchProvider>
</AssistantsProvider>
s
<WrappedSearch
initiallyToggled={toggleSidebar}
searchTypeDefault={searchTypeDefault}
/>
</SearchProvider>
</>
);
}

View File

@ -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>
);
};

View File

@ -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}

View File

@ -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,
};
}

View File

@ -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 llmProviders = (results[7] || []) as LLMProviderDescriptor[];
const foldersResponse = results[8] as Response | null;
const userInputPromptsResponse = results[9] as Response | null;
const tagsResponse = results[5] as Response | null;
const llmProviders = (results[6] || []) as LLMProviderDescriptor[];
const foldersResponse = results[7] 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,
};
}