Skip to content

CC initial draft #5296

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
+243 −87
Closed

CC initial draft #5296

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
242 changes: 164 additions & 78 deletions web/src/app/chat/components/ChatPage.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
useMemo,
useRef,
useState,
memo,
} from "react";
import { usePopup } from "@/components/admin/connectors/Popup";
import { SEARCH_PARAM_NAMES } from "../services/searchParams";
Expand Down Expand Up @@ -92,6 +93,79 @@ import { HumanMessage } from "../message/HumanMessage";
import { AssistantIcon } from "@/components/assistants/AssistantIcon";
import { StarterMessageDisplay } from "./starterMessages/StarterMessageDisplay";

// Memoized wrapper component for AI messages with stable props
const MemoizedAIMessageWrapper = memo(function AIMessageWrapper({
message,
previousMessage,
liveAssistant,
llmModelName,
handleFeedback,
setPresentingDocument,
createRegenerator,
nodeId,
otherMessagesCanSwitchTo,
onMessageSelection,
isLast,
lastMessageRef,
}: {
message: any;
previousMessage: any;
liveAssistant: any;
llmModelName?: string;
handleFeedback: (feedback: FeedbackType, messageId: number) => void;
setPresentingDocument: (doc: MinimalOnyxDocument | null) => void;
createRegenerator: (req: any) => any;
nodeId: number;
otherMessagesCanSwitchTo: number[];
onMessageSelection?: (nodeId: number) => void;
isLast: boolean;
lastMessageRef: React.RefObject<HTMLDivElement>;
}) {
const chatState = useMemo(
() => ({
handleFeedback: (feedback: FeedbackType) =>
handleFeedback(feedback, message.messageId!),
assistant: liveAssistant,
docs: message.documents,
userFiles: [],
citations: message.citations,
setPresentingDocument,
regenerate: createRegenerator({
messageId: message.messageId!,
parentMessage: previousMessage!,
}),
overriddenModel: llmModelName,
}),
[
message.messageId,
message.documents,
message.citations,
liveAssistant,
llmModelName,
previousMessage?.messageId,
handleFeedback,
setPresentingDocument,
createRegenerator,
]
);

return (
<div
className="text-text"
id={`message-${nodeId}`}
ref={isLast ? lastMessageRef : null}
>
<AIMessage
rawPackets={message.packets}
chatState={chatState}
nodeId={nodeId}
otherMessagesCanSwitchTo={otherMessagesCanSwitchTo}
onMessageSelection={onMessageSelection}
/>
</div>
);
});

export function ChatPage({
toggle,
documentSidebarInitialWidth,
Expand All @@ -105,20 +179,20 @@ export function ChatPage({
}) {
// Performance tracking
// Keeping this here in case we need to track down slow renders in the future
// const renderCount = useRef(0);
// renderCount.current++;
// const renderStartTime = performance.now();

// useEffect(() => {
// const renderTime = performance.now() - renderStartTime;
// if (renderTime > 10) {
// console.log(
// `[ChatPage] Slow render #${renderCount.current}: ${renderTime.toFixed(
// 2
// )}ms`
// );
// }
// });
const renderCount = useRef(0);
renderCount.current++;
const renderStartTime = performance.now();

useEffect(() => {
const renderTime = performance.now() - renderStartTime;
if (renderTime > 10) {
console.log(
`[ChatPage] Slow render #${renderCount.current}: ${renderTime.toFixed(
2
)}ms`
);
}
});

const router = useRouter();
const searchParams = useSearchParams();
Expand Down Expand Up @@ -696,22 +770,34 @@ export function ChatPage({
forceSearch?: boolean;
}

function createRegenerator(regenerationRequest: RegenerationRequest) {
// Returns new function that only needs `modelOveride` to be specified when called
return async function (modelOverride: LlmDescriptor) {
return await onSubmit({
message: message,
selectedFiles: selectedFiles,
selectedFolders: selectedFolders,
currentMessageFiles: currentMessageFiles,
useAgentSearch: deepResearchEnabled,
modelOverride,
messageIdToResend: regenerationRequest.parentMessage.messageId,
regenerationRequest,
forceSearch: regenerationRequest.forceSearch,
});
};
}
// Use useCallback to memoize the regenerator creator
// Note: We don't include messageHistory in deps to avoid re-renders
const createRegenerator = useCallback(
(regenerationRequest: RegenerationRequest) => {
// Returns new function that only needs `modelOveride` to be specified when called
return async function (modelOverride: LlmDescriptor) {
// Get the message directly from the regeneration request
return await onSubmit({
message: regenerationRequest.parentMessage.message || "",
selectedFiles: selectedFiles,
selectedFolders: selectedFolders,
currentMessageFiles: currentMessageFiles,
useAgentSearch: deepResearchEnabled,
modelOverride,
messageIdToResend: regenerationRequest.parentMessage.messageId,
regenerationRequest,
forceSearch: regenerationRequest.forceSearch,
});
};
},
[
onSubmit,
selectedFiles,
selectedFolders,
currentMessageFiles,
deepResearchEnabled,
]
);
if (!user) {
redirect("/auth/login");
}
Expand Down Expand Up @@ -773,6 +859,29 @@ export function ChatPage({
setTimeout(() => updateCurrentDocumentSidebarVisible(false), 300);
}, [updateCurrentDocumentSidebarVisible]);

// Memoized callback for handling feedback
const handleFeedback = useCallback(
(feedback: FeedbackType, messageId: number) => {
setCurrentFeedback([feedback, messageId]);
},
[]
);

// Memoized callback for handling message edits
const handleMessageEdit = useCallback(
(editedContent: string, messageId?: number | null) => {
onSubmit({
message: editedContent,
messageIdToResend: messageId || undefined,
selectedFiles: [],
selectedFolders: [],
currentMessageFiles: [],
useAgentSearch: deepResearchEnabled,
});
},
[onSubmit, deepResearchEnabled]
);

// Determine whether to show the centered input (no messages yet)
const showCenteredInput =
messageHistory.length === 0 &&
Expand Down Expand Up @@ -1165,18 +1274,12 @@ export function ChatPage({
content={message.message}
files={message.files}
messageId={message.messageId}
onEdit={(editedContent) => {
onSubmit({
message: editedContent,
messageIdToResend:
message.messageId || undefined,
// TODO: fix
selectedFiles: [],
selectedFolders: [],
currentMessageFiles: [],
useAgentSearch: deepResearchEnabled,
});
}}
onEdit={(editedContent: string) =>
handleMessageEdit(
editedContent,
message.messageId
)
}
otherMessagesCanSwitchTo={
parentMessage?.childrenNodeIds || []
}
Expand Down Expand Up @@ -1208,44 +1311,27 @@ export function ChatPage({
}

return (
<div
className="text-text"
id={`message-${message.nodeId}`}
<MemoizedAIMessageWrapper
key={messageReactComponentKey}
ref={
i == messageHistory.length - 1
? lastMessageRef
: null
message={message}
previousMessage={previousMessage}
liveAssistant={liveAssistant}
llmModelName={
llmManager.currentLlm?.modelName
}
>
<AIMessage
rawPackets={message.packets}
chatState={{
handleFeedback: (feedback) =>
setCurrentFeedback([
feedback,
message.messageId!,
]),
assistant: liveAssistant,
docs: message.documents,
userFiles: [], // TODO: Extract user files from message context
citations: message.citations,
setPresentingDocument:
setPresentingDocument,
regenerate: createRegenerator({
messageId: message.messageId!,
parentMessage: previousMessage!,
}),
overriddenModel:
llmManager.currentLlm?.modelName,
}}
nodeId={message.nodeId}
otherMessagesCanSwitchTo={
parentMessage?.childrenNodeIds || []
}
onMessageSelection={onMessageSelection}
/>
</div>
handleFeedback={handleFeedback}
setPresentingDocument={
setPresentingDocument
}
createRegenerator={createRegenerator}
nodeId={message.nodeId}
otherMessagesCanSwitchTo={
parentMessage?.childrenNodeIds || []
}
onMessageSelection={onMessageSelection}
isLast={i === messageHistory.length - 1}
lastMessageRef={lastMessageRef}
/>
);
}
})}
Expand Down
6 changes: 3 additions & 3 deletions web/src/app/chat/hooks/useChatSessionController.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"use client";

import { useEffect, useRef } from "react";
import { useCallback, useEffect, useRef } from "react";
import { ReadonlyURLSearchParams, useRouter } from "next/navigation";
import {
nameChatSession,
Expand Down Expand Up @@ -305,7 +305,7 @@ export function useChatSessionController({
// This effect should only run when existingChatSessionId or persona ID changes
]);

const onMessageSelection = (nodeId: number) => {
const onMessageSelection = useCallback((nodeId: number) => {
updateCurrentSelectedNodeForDocDisplay(nodeId);
const currentMessageTree = useChatSessionStore
.getState()
Expand All @@ -330,7 +330,7 @@ export function useChatSessionController({
console.error("Message has no messageId", nodeId);
}
}
};
}, []);

return {
onMessageSelection,
Expand Down
39 changes: 37 additions & 2 deletions web/src/app/chat/message/HumanMessage.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"use client";

import { FiEdit2 } from "react-icons/fi";
import React, { useEffect, useRef, useState } from "react";
import React, { useEffect, useRef, useState, memo } from "react";
import { MinimalOnyxDocument } from "@/lib/search/interfaces";
import { ChatFileType, FileDescriptor } from "@/app/chat/interfaces";
import { Hoverable, HoverableIcon } from "@/components/Hoverable";
Expand Down Expand Up @@ -103,7 +103,7 @@ function FileDisplay({
);
}

export const HumanMessage = ({
const HumanMessageComponent = ({
content,
files,
messageId,
Expand All @@ -128,6 +128,8 @@ export const HumanMessage = ({
}) => {
const textareaRef = useRef<HTMLTextAreaElement>(null);

console.log("Rendering HumanMessage");

const [isHovered, setIsHovered] = useState(false);
const [isEditing, setIsEditing] = useState(false);
const [editedContent, setEditedContent] = useState(content);
Expand Down Expand Up @@ -406,3 +408,36 @@ export const HumanMessage = ({
</div>
);
};

// Custom comparison function to prevent unnecessary re-renders
const areEqual = (prevProps: any, nextProps: any) => {
// Only re-render if the actual message data has changed
if (prevProps.content !== nextProps.content) {
console.log("HumanMessage re-render: content changed");
return false;
}
if (prevProps.messageId !== nextProps.messageId) {
console.log("HumanMessage re-render: messageId changed");
return false;
}
if (
prevProps.disableSwitchingForStreaming !==
nextProps.disableSwitchingForStreaming
) {
console.log("HumanMessage re-render: disableSwitchingForStreaming changed");
return false;
}

// Check files array - only if length changed
if (prevProps.files?.length !== nextProps.files?.length) {
console.log("HumanMessage re-render: files length changed");
return false;
}

// Ignore all other prop changes including functions and arrays
// that are just reference changes
return true;
};

export const HumanMessage = memo(HumanMessageComponent, areEqual);
HumanMessage.displayName = "HumanMessage";
Loading
Loading