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,21 +10,109 @@ interface CodeBlockProps {
   [key: string]: any;
 }
 
-export function CodeBlock({
+export const CodeBlock = memo(function CodeBlock({
   className = "",
   children,
   content,
   ...props
 }: CodeBlockProps) {
-  const language = className
-    .split(" ")
-    .filter((cls) => cls.startsWith("language-"))
-    .map((cls) => cls.replace("language-", ""))
-    .join(" ");
   const [copied, setCopied] = useState(false);
 
+  const language = useMemo(() => {
+    return className
+      .split(" ")
+      .filter((cls) => cls.startsWith("language-"))
+      .map((cls) => cls.replace("language-", ""))
+      .join(" ");
+  }, [className]);
+
+  const codeText = useMemo(() => {
+    let codeText: string | null = null;
+    if (
+      props.node?.position?.start?.offset &&
+      props.node?.position?.end?.offset
+    ) {
+      codeText = content.slice(
+        props.node.position.start.offset,
+        props.node.position.end.offset
+      );
+      codeText = codeText.trim();
+
+      // Find the last occurrence of closing backticks
+      const lastBackticksIndex = codeText.lastIndexOf("```");
+      if (lastBackticksIndex !== -1) {
+        codeText = codeText.slice(0, lastBackticksIndex + 3);
+      }
+
+      // Remove the language declaration and trailing backticks
+      const codeLines = codeText.split("\n");
+      if (
+        codeLines.length > 1 &&
+        (codeLines[0].startsWith("```") ||
+          codeLines[0].trim().startsWith("```"))
+      ) {
+        codeLines.shift(); // Remove the first line with the language declaration
+        if (
+          codeLines[codeLines.length - 1] === "```" ||
+          codeLines[codeLines.length - 1]?.trim() === "```"
+        ) {
+          codeLines.pop(); // Remove the last line with the trailing backticks
+        }
+
+        const minIndent = codeLines
+          .filter((line) => line.trim().length > 0)
+          .reduce((min, line) => {
+            const match = line.match(/^\s*/);
+            return Math.min(min, match ? match[0].length : 0);
+          }, Infinity);
+
+        const formattedCodeLines = codeLines.map((line) =>
+          line.slice(minIndent)
+        );
+        codeText = formattedCodeLines.join("\n");
+      }
+    }
+
+    // handle unknown languages. They won't have a `node.position.start.offset`
+    if (!codeText) {
+      const findTextNode = (node: any): string | null => {
+        if (node.type === "text") {
+          return node.value;
+        }
+        let finalResult = "";
+        if (node.children) {
+          for (const child of node.children) {
+            const result = findTextNode(child);
+            if (result) {
+              finalResult += result;
+            }
+          }
+        }
+        return finalResult;
+      };
+
+      codeText = findTextNode(props.node);
+    }
+
+    return codeText;
+  }, [content, props.node]);
+
+  const handleCopy = useCallback(
+    (event: React.MouseEvent) => {
+      event.preventDefault();
+      if (!codeText) {
+        return;
+      }
+
+      navigator.clipboard.writeText(codeText).then(() => {
+        setCopied(true);
+        setTimeout(() => setCopied(false), 2000);
+      });
+    },
+    [codeText]
+  );
+
   if (!language) {
-    // this is the case of a single "`" e.g. `hi`
     if (typeof children === "string") {
       return <code className={className}>{children}</code>;
     }
@@ -39,82 +126,6 @@ export function CodeBlock({
     );
   }
 
-  let codeText: string | null = null;
-  if (
-    props.node?.position?.start?.offset &&
-    props.node?.position?.end?.offset
-  ) {
-    codeText = content.slice(
-      props.node.position.start.offset,
-      props.node.position.end.offset
-    );
-    codeText = codeText.trim();
-
-    // Find the last occurrence of closing backticks
-    const lastBackticksIndex = codeText.lastIndexOf("```");
-    if (lastBackticksIndex !== -1) {
-      codeText = codeText.slice(0, lastBackticksIndex + 3);
-    }
-
-    // Remove the language declaration and trailing backticks
-    const codeLines = codeText.split("\n");
-    if (
-      codeLines.length > 1 &&
-      (codeLines[0].startsWith("```") || codeLines[0].trim().startsWith("```"))
-    ) {
-      codeLines.shift(); // Remove the first line with the language declaration
-      if (
-        codeLines[codeLines.length - 1] === "```" ||
-        codeLines[codeLines.length - 1]?.trim() === "```"
-      ) {
-        codeLines.pop(); // Remove the last line with the trailing backticks
-      }
-
-      const minIndent = codeLines
-        .filter((line) => line.trim().length > 0)
-        .reduce((min, line) => {
-          const match = line.match(/^\s*/);
-          return Math.min(min, match ? match[0].length : 0);
-        }, Infinity);
-
-      const formattedCodeLines = codeLines.map((line) => line.slice(minIndent));
-      codeText = formattedCodeLines.join("\n");
-    }
-  }
-
-  // handle unknown languages. They won't have a `node.position.start.offset`
-  if (!codeText) {
-    const findTextNode = (node: any): string | null => {
-      if (node.type === "text") {
-        return node.value;
-      }
-      let finalResult = "";
-      if (node.children) {
-        for (const child of node.children) {
-          const result = findTextNode(child);
-          if (result) {
-            finalResult += result;
-          }
-        }
-      }
-      return finalResult;
-    };
-
-    codeText = findTextNode(props.node);
-  }
-
-  const handleCopy = (event: React.MouseEvent) => {
-    event.preventDefault();
-    if (!codeText) {
-      return;
-    }
-
-    navigator.clipboard.writeText(codeText).then(() => {
-      setCopied(true);
-      setTimeout(() => setCopied(false), 2000); // Reset copy status after 2 seconds
-    });
-  };
-
   return (
     <div className="overflow-x-hidden">
       <div className="flex mx-3 py-2 text-xs">
@@ -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} />