Memoize AI message component (#2483)

* memoize AI message component

* rename memoized file

* remove "zz"

* update name

* memoize for coverage

* add display name

web/src/app/chat/message/CodeBlock.tsx

@@ -1,5 +1,4 @@
-import React from "react";
+import React, { useState, ReactNode, useCallback, useMemo, memo } from "react";
-import { useState, ReactNode } from "react";
 import { FiCheck, FiCopy } from "react-icons/fi";
 
 const CODE_BLOCK_PADDING_TYPE = { padding: "1rem" };
@@ -11,34 +10,23 @@ interface CodeBlockProps {
   [key: string]: any;
 }
 
-export function CodeBlock({
+export const CodeBlock = memo(function CodeBlock({
   className = "",
   children,
   content,
   ...props
 }: CodeBlockProps) {
-  const language = className
+  const [copied, setCopied] = useState(false);
+
+  const language = useMemo(() => {
+    return className
       .split(" ")
       .filter((cls) => cls.startsWith("language-"))
       .map((cls) => cls.replace("language-", ""))
       .join(" ");
-  const [copied, setCopied] = useState(false);
+  }, [className]);
-
-  if (!language) {
-    // this is the case of a single "`" e.g. `hi`
-    if (typeof children === "string") {
-      return <code className={className}>{children}</code>;
-    }
-
-    return (
-      <pre style={CODE_BLOCK_PADDING_TYPE}>
-        <code {...props} className={`text-sm ${className}`}>
-          {children}
-        </code>
-      </pre>
-    );
-  }
 
+  const codeText = useMemo(() => {
     let codeText: string | null = null;
     if (
       props.node?.position?.start?.offset &&
@@ -60,7 +48,8 @@ export function CodeBlock({
       const codeLines = codeText.split("\n");
       if (
         codeLines.length > 1 &&
-      (codeLines[0].startsWith("```") || codeLines[0].trim().startsWith("```"))
+        (codeLines[0].startsWith("```") ||
+          codeLines[0].trim().startsWith("```"))
       ) {
         codeLines.shift(); // Remove the first line with the language declaration
         if (
@@ -77,7 +66,9 @@ export function CodeBlock({
             return Math.min(min, match ? match[0].length : 0);
           }, Infinity);
 
-      const formattedCodeLines = codeLines.map((line) => line.slice(minIndent));
+        const formattedCodeLines = codeLines.map((line) =>
+          line.slice(minIndent)
+        );
         codeText = formattedCodeLines.join("\n");
       }
     }
@@ -103,7 +94,11 @@ export function CodeBlock({
       codeText = findTextNode(props.node);
     }
 
-  const handleCopy = (event: React.MouseEvent) => {
+    return codeText;
+  }, [content, props.node]);
+
+  const handleCopy = useCallback(
+    (event: React.MouseEvent) => {
       event.preventDefault();
       if (!codeText) {
         return;
@@ -111,9 +106,25 @@ export function CodeBlock({
 
       navigator.clipboard.writeText(codeText).then(() => {
         setCopied(true);
-      setTimeout(() => setCopied(false), 2000); // Reset copy status after 2 seconds
+        setTimeout(() => setCopied(false), 2000);
       });
-  };
+    },
+    [codeText]
+  );
+
+  if (!language) {
+    if (typeof children === "string") {
+      return <code className={className}>{children}</code>;
+    }
+
+    return (
+      <pre style={CODE_BLOCK_PADDING_TYPE}>
+        <code {...props} className={`text-sm ${className}`}>
+          {children}
+        </code>
+      </pre>
+    );
+  }
 
   return (
     <div className="overflow-x-hidden">
@@ -143,4 +154,4 @@ export function CodeBlock({
       </pre>
     </div>
   );
-}
+});

web/src/app/chat/message/MemoizedTextComponents.tsx

@@ -0,0 +1,33 @@
+import { Citation } from "@/components/search/results/Citation";
+import React, { memo } from "react";
+
+export const MemoizedLink = memo((props: any) => {
+  const { node, ...rest } = props;
+  const value = rest.children;
+
+  if (value?.toString().startsWith("*")) {
+    return (
+      <div className="flex-none bg-background-800 inline-block rounded-full h-3 w-3 ml-2" />
+    );
+  } else if (value?.toString().startsWith("[")) {
+    return <Citation link={rest?.href}>{rest.children}</Citation>;
+  } else {
+    return (
+      <a
+        onMouseDown={() =>
+          rest.href ? window.open(rest.href, "_blank") : undefined
+        }
+        className="cursor-pointer text-link hover:text-link-hover"
+      >
+        {rest.children}
+      </a>
+    );
+  }
+});
+
+export const MemoizedParagraph = memo(({ node, ...props }: any) => (
+  <p {...props} className="text-default" />
+));
+
+MemoizedLink.displayName = "MemoizedLink";
+MemoizedParagraph.displayName = "MemoizedParagraph";

web/src/app/chat/message/Messages.tsx

@@ -8,14 +8,7 @@ import {
   FiGlobe,
 } from "react-icons/fi";
 import { FeedbackType } from "../types";
-import {
+import React, { useContext, useEffect, useMemo, useRef, useState } from "react";
-  Dispatch,
-  SetStateAction,
-  useContext,
-  useEffect,
-  useRef,
-  useState,
-} from "react";
 import ReactMarkdown from "react-markdown";
 import {
   DanswerDocument,
@@ -46,12 +39,7 @@ import { AssistantIcon } from "@/components/assistants/AssistantIcon";
 import { Citation } from "@/components/search/results/Citation";
 import { DocumentMetadataBlock } from "@/components/search/DocumentDisplay";
 
-import {
+import { LikeFeedback, DislikeFeedback } from "@/components/icons/icons";
-  ThumbsUpIcon,
-  ThumbsDownIcon,
-  LikeFeedback,
-  DislikeFeedback,
-} from "@/components/icons/icons";
 import {
   CustomTooltip,
   TooltipGroup,
@@ -65,6 +53,7 @@ import GeneratingImageDisplay from "../tools/GeneratingImageDisplay";
 import RegenerateOption from "../RegenerateOption";
 import { LlmOverride } from "@/lib/hooks";
 import { ContinueGenerating } from "./ContinueMessage";
+import { MemoizedLink, MemoizedParagraph } from "./MemoizedTextComponents";
 
 const TOOLS_WITH_CUSTOM_HANDLING = [
   SEARCH_TOOL_NAME,
@@ -367,41 +356,8 @@ export const AIMessage = ({
                               key={messageId}
                               className="prose max-w-full text-base"
                               components={{
-                                a: (props) => {
+                                a: MemoizedLink,
-                                  const { node, ...rest } = props;
+                                p: MemoizedParagraph,
-                                  const value = rest.children;
-
-                                  if (value?.toString().startsWith("*")) {
-                                    return (
-                                      <div className="flex-none bg-background-800 inline-block rounded-full h-3 w-3 ml-2" />
-                                    );
-                                  } else if (
-                                    value?.toString().startsWith("[")
-                                  ) {
-                                    // for some reason <a> tags cause the onClick to not apply
-                                    // and the links are unclickable
-                                    // TODO: fix the fact that you have to double click to follow link
-                                    // for the first link
-                                    return (
-                                      <Citation link={rest?.href}>
-                                        {rest.children}
-                                      </Citation>
-                                    );
-                                  } else {
-                                    return (
-                                      <a
-                                        onMouseDown={() =>
-                                          rest.href
-                                            ? window.open(rest.href, "_blank")
-                                            : undefined
-                                        }
-                                        className="cursor-pointer text-link hover:text-link-hover"
-                                      >
-                                        {rest.children}
-                                      </a>
-                                    );
-                                  }
-                                },
                                 code: (props) => (
                                   <CodeBlock
                                     className="w-full"
@@ -409,9 +365,6 @@ export const AIMessage = ({
                                     content={content as string}
                                   />
                                 ),
-                                p: ({ node, ...props }) => (
-                                  <p {...props} className="text-default" />
-                                ),
                               }}
                               remarkPlugins={[remarkGfm]}
                               rehypePlugins={[

web/src/components/chat_search/MinimalMarkdown.tsx

@@ -1,4 +1,8 @@
 import { CodeBlock } from "@/app/chat/message/CodeBlock";
+import {
+  MemoizedLink,
+  MemoizedParagraph,
+} from "@/app/chat/message/MemoizedTextComponents";
 import React from "react";
 import ReactMarkdown from "react-markdown";
 import remarkGfm from "remark-gfm";
@@ -18,17 +22,8 @@ export const MinimalMarkdown: React.FC<MinimalMarkdownProps> = ({
     <ReactMarkdown
       className={`w-full text-wrap break-word ${className}`}
       components={{
-        a: ({ node, ...props }) => (
+        a: MemoizedLink,
-          <a
+        p: MemoizedParagraph,
-            {...props}
-            className="text-sm text-link hover:text-link-hover"
-            target="_blank"
-            rel="noopener noreferrer"
-          />
-        ),
-        p: ({ node, ...props }) => (
-          <p {...props} className="text-wrap break-word text-sm m-0 w-full" />
-        ),
         code: useCodeBlock
           ? (props) => (
               <CodeBlock className="w-full" {...props} content={content} />