Skip to content

Commit 0fd268d

Browse files
WevesWeves
authored
fix: message render performance (#5297)
* Separate out message display into it's own component * Memoize AIMessage * Cleanup * Remove log * Address greptile/cubic comments
1 parent f345da7 commit 0fd268d

File tree

8 files changed

+585
-257
lines changed

8 files changed

+585
-257
lines changed

‎web/src/app/chat/components/ChatPage.tsx‎

Lines changed: 45 additions & 199 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,7 @@ import {
1717
} from "react";
1818
import { usePopup } from "@/components/admin/connectors/Popup";
1919
import { SEARCH_PARAM_NAMES } from "../services/searchParams";
20-
import {
21-
LlmDescriptor,
22-
useFederatedConnectors,
23-
useFilters,
24-
useLlmManager,
25-
} from "@/lib/hooks";
20+
import { useFederatedConnectors, useFilters, useLlmManager } from "@/lib/hooks";
2621
import { FeedbackType } from "@/app/chat/interfaces";
2722
import { OnyxInitializingLoader } from "@/components/OnyxInitializingLoader";
2823
import { FeedbackModal } from "./modal/FeedbackModal";
@@ -86,11 +81,10 @@ import {
8681
useChatSessionSharedStatus,
8782
useHasSentLocalUserMessage,
8883
} from "../stores/useChatSessionStore";
89-
import { AIMessage } from "../message/messageComponents/AIMessage";
9084
import { FederatedOAuthModal } from "@/components/chat/FederatedOAuthModal";
91-
import { HumanMessage } from "../message/HumanMessage";
9285
import { AssistantIcon } from "@/components/assistants/AssistantIcon";
9386
import { StarterMessageDisplay } from "./starterMessages/StarterMessageDisplay";
87+
import { MessagesDisplay } from "./MessagesDisplay";
9488

9589
export function ChatPage({
9690
toggle,
@@ -360,15 +354,15 @@ export function ChatPage({
360354
}, 100);
361355
};
362356

363-
const resetInputBar = () => {
357+
const resetInputBar = useCallback(() => {
364358
setMessage("");
365359
setCurrentMessageFiles([]);
366360
if (endPaddingRef.current) {
367361
endPaddingRef.current.style.height = `95px`;
368362
}
369-
};
363+
}, [setMessage, setCurrentMessageFiles]);
370364

371-
const clientScrollToBottom = (fast?: boolean) => {
365+
const clientScrollToBottom = useCallback((fast?: boolean) => {
372366
waitForScrollRef.current = true;
373367

374368
setTimeout(() => {
@@ -389,15 +383,17 @@ export function ChatPage({
389383
});
390384

391385
if (chatSessionIdRef.current) {
392-
updateHasPerformedInitialScroll(chatSessionIdRef.current, true);
386+
useChatSessionStore
387+
.getState()
388+
.updateHasPerformedInitialScroll(chatSessionIdRef.current, true);
393389
}
394390
}, 50);
395391

396392
// Reset waitForScrollRef after 1.5 seconds
397393
setTimeout(() => {
398394
waitForScrollRef.current = false;
399395
}, 1500);
400-
};
396+
}, []);
401397

402398
const debounceNumber = 100; // time for debouncing
403399

@@ -470,7 +466,6 @@ export function ChatPage({
470466
);
471467

472468
// Access chat state directly from the store
473-
const beforeZustandTime = performance.now();
474469
const currentChatState = useCurrentChatState();
475470
const chatSessionId = useChatSessionStore((state) => state.currentSessionId);
476471
const submittedMessage = useSubmittedMessage();
@@ -690,28 +685,6 @@ export function ChatPage({
690685
[setIsChatSearchModalOpen]
691686
);
692687

693-
interface RegenerationRequest {
694-
messageId: number;
695-
parentMessage: Message;
696-
forceSearch?: boolean;
697-
}
698-
699-
function createRegenerator(regenerationRequest: RegenerationRequest) {
700-
// Returns new function that only needs `modelOveride` to be specified when called
701-
return async function (modelOverride: LlmDescriptor) {
702-
return await onSubmit({
703-
message: message,
704-
selectedFiles: selectedFiles,
705-
selectedFolders: selectedFolders,
706-
currentMessageFiles: currentMessageFiles,
707-
useAgentSearch: deepResearchEnabled,
708-
modelOverride,
709-
messageIdToResend: regenerationRequest.parentMessage.messageId,
710-
regenerationRequest,
711-
forceSearch: regenerationRequest.forceSearch,
712-
});
713-
};
714-
}
715688
if (!user) {
716689
redirect("/auth/login");
717690
}
@@ -853,6 +826,9 @@ export function ChatPage({
853826
onOutsideClick={() => updateCurrentDocumentSidebarVisible(false)}
854827
title="Sources"
855828
>
829+
{/* IMPORTANT: this is a memoized component, and it's very important
830+
for performance reasons that this stays true. MAKE SURE that all function
831+
props are wrapped in useCallback. */}
856832
<DocumentResults
857833
setPresentingDocument={setPresentingDocument}
858834
modal={true}
@@ -1008,6 +984,9 @@ export function ChatPage({
1008984
}
1009985
`}
1010986
>
987+
{/* IMPORTANT: this is a memoized component, and it's very important
988+
for performance reasons that this stays true. MAKE SURE that all function
989+
props are wrapped in useCallback. */}
1011990
<DocumentResults
1012991
setPresentingDocument={setPresentingDocument}
1013992
modal={false}
@@ -1114,171 +1093,38 @@ export function ChatPage({
11141093
</div>
11151094
)}
11161095

1117-
{messageHistory.length === 0 &&
1118-
!isFetchingChatMessages &&
1119-
!loadingError &&
1120-
!submittedMessage &&
1121-
null}
1122-
<div
1123-
style={{ overflowAnchor: "none" }}
1124-
key={chatSessionId}
1125-
className={
1126-
(hasPerformedInitialScroll ? "" : " hidden ") +
1127-
"desktop:-ml-4 w-full mx-auto " +
1128-
"absolute mobile:top-0 desktop:top-0 left-0 " +
1129-
(settings?.enterpriseSettings
1130-
?.two_lines_for_chat_header
1131-
? "pt-20 "
1132-
: "pt-4 ")
1096+
<MessagesDisplay
1097+
messageHistory={messageHistory}
1098+
completeMessageTree={completeMessageTree}
1099+
liveAssistant={liveAssistant}
1100+
llmManager={llmManager}
1101+
deepResearchEnabled={deepResearchEnabled}
1102+
selectedFiles={selectedFiles}
1103+
selectedFolders={selectedFolders}
1104+
currentMessageFiles={currentMessageFiles}
1105+
setPresentingDocument={setPresentingDocument}
1106+
setCurrentFeedback={setCurrentFeedback}
1107+
onSubmit={onSubmit}
1108+
onMessageSelection={onMessageSelection}
1109+
stopGenerating={stopGenerating}
1110+
uncaughtError={uncaughtError}
1111+
loadingError={loadingError}
1112+
handleResubmitLastMessage={
1113+
handleResubmitLastMessage
11331114
}
1134-
// NOTE: temporarily removing this to fix the scroll bug
1135-
// (hasPerformedInitialScroll ? "" : "invisible")
1136-
>
1137-
{messageHistory.map((message, i) => {
1138-
const messageTree = completeMessageTree;
1139-
1140-
const messageReactComponentKey = `message-${message.nodeId}`;
1141-
const parentMessage = message.parentNodeId
1142-
? messageTree?.get(message.parentNodeId)
1143-
: null;
1144-
if (message.type === "user") {
1145-
const nextMessage =
1146-
messageHistory.length > i + 1
1147-
? messageHistory[i + 1]
1148-
: null;
1149-
1150-
return (
1151-
<div
1152-
id={messageReactComponentKey}
1153-
key={messageReactComponentKey}
1154-
>
1155-
<HumanMessage
1156-
setPresentingDocument={
1157-
setPresentingDocument
1158-
}
1159-
disableSwitchingForStreaming={
1160-
(nextMessage &&
1161-
nextMessage.is_generating) ||
1162-
false
1163-
}
1164-
stopGenerating={stopGenerating}
1165-
content={message.message}
1166-
files={message.files}
1167-
messageId={message.messageId}
1168-
onEdit={(editedContent) => {
1169-
onSubmit({
1170-
message: editedContent,
1171-
messageIdToResend:
1172-
message.messageId || undefined,
1173-
// TODO: fix
1174-
selectedFiles: [],
1175-
selectedFolders: [],
1176-
currentMessageFiles: [],
1177-
useAgentSearch: deepResearchEnabled,
1178-
});
1179-
}}
1180-
otherMessagesCanSwitchTo={
1181-
parentMessage?.childrenNodeIds || []
1182-
}
1183-
onMessageSelection={onMessageSelection}
1184-
/>
1185-
</div>
1186-
);
1187-
} else if (message.type === "assistant") {
1188-
const previousMessage =
1189-
i !== 0 ? messageHistory[i - 1] : null;
1190-
1191-
if (
1192-
(uncaughtError || loadingError) &&
1193-
i === messageHistory.length - 1
1194-
) {
1195-
return (
1196-
<div
1197-
key={`error-${message.nodeId}`}
1198-
className="max-w-message-max mx-auto"
1199-
>
1200-
<ErrorBanner
1201-
resubmit={handleResubmitLastMessage}
1202-
error={
1203-
uncaughtError || loadingError || ""
1204-
}
1205-
/>
1206-
</div>
1207-
);
1208-
}
1209-
1210-
return (
1211-
<div
1212-
className="text-text"
1213-
id={`message-${message.nodeId}`}
1214-
key={messageReactComponentKey}
1215-
ref={
1216-
i == messageHistory.length - 1
1217-
? lastMessageRef
1218-
: null
1219-
}
1220-
>
1221-
<AIMessage
1222-
rawPackets={message.packets}
1223-
chatState={{
1224-
handleFeedback: (feedback) =>
1225-
setCurrentFeedback([
1226-
feedback,
1227-
message.messageId!,
1228-
]),
1229-
assistant: liveAssistant,
1230-
docs: message.documents,
1231-
userFiles: [], // TODO: Extract user files from message context
1232-
citations: message.citations,
1233-
setPresentingDocument:
1234-
setPresentingDocument,
1235-
regenerate: createRegenerator({
1236-
messageId: message.messageId!,
1237-
parentMessage: previousMessage!,
1238-
}),
1239-
overriddenModel:
1240-
llmManager.currentLlm?.modelName,
1241-
}}
1242-
nodeId={message.nodeId}
1243-
otherMessagesCanSwitchTo={
1244-
parentMessage?.childrenNodeIds || []
1245-
}
1246-
onMessageSelection={onMessageSelection}
1247-
/>
1248-
</div>
1249-
);
1250-
}
1251-
})}
1252-
1253-
{((uncaughtError || loadingError) &&
1254-
messageHistory[messageHistory.length - 1]
1255-
?.type === "user") ||
1256-
(messageHistory[messageHistory.length - 1]
1257-
?.type === "error" && (
1258-
<div className="max-w-message-max mx-auto">
1259-
<ErrorBanner
1260-
resubmit={handleResubmitLastMessage}
1261-
error={uncaughtError || loadingError || ""}
1262-
/>
1263-
</div>
1264-
))}
1265-
1266-
{messageHistory.length > 0 && (
1267-
<div
1268-
style={{
1269-
height: !autoScrollEnabled
1270-
? getContainerHeight()
1271-
: undefined,
1272-
}}
1273-
/>
1274-
)}
1275-
1276-
{/* Some padding at the bottom so the search bar has space at the bottom to not cover the last message*/}
1277-
<div ref={endPaddingRef} className="h-[95px]" />
1278-
1279-
<div ref={endDivRef} />
1280-
</div>
1115+
autoScrollEnabled={autoScrollEnabled}
1116+
getContainerHeight={getContainerHeight}
1117+
lastMessageRef={lastMessageRef}
1118+
endPaddingRef={endPaddingRef}
1119+
endDivRef={endDivRef}
1120+
hasPerformedInitialScroll={
1121+
hasPerformedInitialScroll
1122+
}
1123+
chatSessionId={chatSessionId}
1124+
enterpriseSettings={enterpriseSettings}
1125+
/>
12811126
</div>
1127+
12821128
<div
12831129
ref={inputRef}
12841130
className={`absolute pointer-events-none z-10 w-full ${

0 commit comments

Comments
 (0)