fix: Math \boldsymbol is not rendered (#184)

Co-authored-by: Chris Chudzicki <[email protected]>

Author: jonkafton

src/bundles/AiDrawer/AiDrawerManager.tsx

@@ -2,6 +2,7 @@ import * as React from "react"
 import { useEffect, useState } from "react"
 import { AiDrawer } from "./AiDrawer"
 import type { AiDrawerProps, AiDrawerSettings } from "./AiDrawer"
+import { contentHash } from "../../utils/string"
 
 type AiDrawerInitMessage = {
   type: "smoot-design::ai-drawer-open" | "smoot-design::tutor-drawer-open" // ("smoot-design::tutor-drawer-open" is legacy)
@@ -16,12 +17,7 @@ type AiDrawerInitMessage = {
 
 const hashPayload = (payload: AiDrawerInitMessage["payload"]) => {
   const str = JSON.stringify(payload)
-  let hash = 5381
-  for (let i = 0; i < str.length; i++) {
-    hash = (hash << 5) + hash + str.charCodeAt(i)
-    hash = hash & hash
-  }
-  return Math.abs(hash).toString(36)
+  return contentHash(str)
 }
 
 type AiDrawerManagerProps = {

src/components/AiChat/AiChat.test.tsx

@@ -2,7 +2,8 @@
 /* eslint-disable testing-library/await-async-utils */
 import { render, screen, waitFor } from "@testing-library/react"
 import user from "@testing-library/user-event"
-import { AiChat, replaceMathjax } from "./AiChat"
+import { AiChat } from "./AiChat"
+import { replaceMathjax } from "./Markdown"
 import { ThemeProvider } from "../ThemeProvider/ThemeProvider"
 import * as React from "react"
 import { AiChatProps } from "./types"

src/components/AiChat/AiChat.tsx

@@ -21,16 +21,37 @@ import EllipsisIcon from "./EllipsisIcon"
 import { SimpleSelectField } from "../SimpleSelect/SimpleSelect"
 import { useFetch } from "./utils"
 import { SelectChangeEvent } from "@mui/material/Select"
+import type { MathJax3Config } from "better-react-mathjax"
 import { MathJaxContext } from "better-react-mathjax"
+import deepmerge from "@mui/utils/deepmerge"
 
 const ConditionalMathJaxWrapper: React.FC<{
   useMathJax: boolean
+  config?: MathJax3Config
   children: React.ReactNode
-}> = ({ useMathJax, children }) => {
+}> = ({ useMathJax, config = {}, children }) => {
   if (!useMathJax) {
     return <>{children}</>
   }
-  return <MathJaxContext>{children}</MathJaxContext>
+
+  return (
+    <MathJaxContext
+      config={deepmerge(
+        {
+          startup: {
+            typeset: false,
+          },
+          loader: { load: ["[tex]/boldsymbol"] },
+          tex: {
+            packages: { "[+]": ["boldsymbol"] },
+          },
+        },
+        config,
+      )}
+    >
+      {children}
+    </MathJaxContext>
+  )
 }
 
 const classes = {
@@ -251,13 +272,13 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
   scrollElement,
   ref,
   useMathJax = false,
+  mathJaxConfig,
   onSubmit,
   problemSetListUrl,
   problemSetInitialMessages,
   problemSetEmptyMessages,
   ...others // Could contain data attributes
 }) => {
-  const containerRef = useRef<HTMLDivElement>(null)
   const messagesContainerRef = useRef<HTMLDivElement>(null)
   const chatScreenRef = useRef<HTMLDivElement>(null)
   const promptInputRef = useRef<HTMLDivElement>(null)
@@ -291,7 +312,9 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
   const [showEntryScreen, setShowEntryScreen] = useState(entryScreenEnabled)
   useEffect(() => {
     if (!showEntryScreen) {
-      promptInputRef.current?.querySelector("input")?.focus()
+      promptInputRef.current
+        ?.querySelector("input")
+        ?.focus({ preventScroll: true })
     }
   }, [showEntryScreen])
 
@@ -318,7 +341,7 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
         ])
       }
     }
-  }, [problemSetListResponse])
+  }, [problemSetListResponse, problemSetEmptyMessages, setMessages])
 
   useEffect(() => {
     if (
@@ -366,7 +389,7 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
   const externalScroll = !!scrollElement
 
   return (
-    <Container className={className} ref={containerRef}>
+    <Container className={className}>
       {showEntryScreen ? (
         <EntryScreen
           className={classes.entryScreenContainer}
@@ -420,41 +443,37 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
                 ) : null
               }
             />
-            <ConditionalMathJaxWrapper useMathJax={useMathJax}>
+            <ConditionalMathJaxWrapper
+              useMathJax={useMathJax}
+              config={mathJaxConfig}
+            >
               <MessagesContainer
                 className={classes.messagesContainer}
                 externalScroll={externalScroll}
                 ref={messagesContainerRef}
               >
-                {messages.map((m: Message, i) => {
-                  // Our Markdown+Mathjax has issues when rendering streaming display math
-                  // Force a re-render of the last (streaming) message when it's done loading.
-                  const key =
-                    i === messages.length - 1 && isLoading
-                      ? `isLoading-${m.id}`
-                      : m.id
+                {messages.map((message: Message, index: number) => {
                   return (
                     <MessageRow
-                      key={key}
-                      data-chat-role={m.role}
+                      key={index}
+                      data-chat-role={message.role}
                       className={classNames(classes.messageRow, {
-                        [classes.messageRowUser]: m.role === "user",
-                        [classes.messageRowAssistant]: m.role === "assistant",
+                        [classes.messageRowUser]: message.role === "user",
+                        [classes.messageRowAssistant]:
+                          message.role === "assistant",
                       })}
                     >
                       <Message className={classes.message}>
-                        <VisuallyHidden as={m.role === "user" ? "h5" : "h6"}>
-                          {m.role === "user"
+                        <VisuallyHidden
+                          as={message.role === "user" ? "h5" : "h6"}
+                        >
+                          {message.role === "user"
                             ? "You said: "
                             : "Assistant said: "}
                         </VisuallyHidden>
-                        {useMathJax ? (
-                          <Markdown enableMathjax={true}>
-                            {replaceMathjax(m.content)}
-                          </Markdown>
-                        ) : (
-                          <Markdown>{m.content}</Markdown>
-                        )}
+                        <Markdown useMathJax={useMathJax}>
+                          {message.content}
+                        </Markdown>
                       </Message>
                     </MessageRow>
                   )
@@ -583,21 +602,4 @@ const AiChat: FC<AiChatProps> = ({
   )
 }
 
-// react-markdown expects Mathjax delimiters to be $...$ or $$...$$
-// the prompt for the tutorbot asks for Mathjax tags with $ format but
-// the LLM does not get it right all the time
-// this function replaces the Mathjax tags with the correct format
-// eventually we will probably be able to remove this as LLMs get better
-function replaceMathjax(inputString: string): string {
-  // Replace instances of \(...\) and \[...\] Mathjax tags with $...$
-  // and $$...$$ tags.
-  const INLINE_MATH_REGEX = /\\\((.*?)\\\)/g
-  const DISPLAY_MATH_REGEX = /\\\[(([\s\S]*?))\\\]/g
-  inputString = inputString.replace(
-    INLINE_MATH_REGEX,
-    (_match, p1) => `$${p1}$`,
-  )
-  return inputString.replace(DISPLAY_MATH_REGEX, (_match, p1) => `$$${p1}$$`)
-}
-
-export { AiChatDisplay, AiChat, replaceMathjax }
+export { AiChatDisplay, AiChat }

src/components/AiChat/AiChatMarkdown.stories.tsx

@@ -5,10 +5,10 @@ import styled from "@emotion/styled"
 import { handlers } from "./test-utils/api"
 
 const TEST_API_STREAMING = "http://localhost:4567/streaming"
+const TEST_API_STREAMING_MATH = "http://localhost:4567/streaming-math"
 
 const Container = styled.div({
   width: "100%",
-  height: "500px",
 })
 
 const meta: Meta<typeof AiChat> = {
@@ -20,7 +20,7 @@ const meta: Meta<typeof AiChat> = {
   render: (args) => <AiChat {...args} />,
   decorators: (Story, context) => {
     return (
-      <Container>
+      <Container style={{ height: context.parameters.height || "500px" }}>
         <Story key={String(context.args.entryScreenEnabled)} />
       </Container>
     )
@@ -122,6 +122,12 @@ def f(x):
   },
 }
 
+/**
+ * Shows MathJax rendering of inline and block math.
+ *
+ * We set `startup.typeset` to false in the MathJax config as by default MathJax will typeset math anywhere on the page and outside of our components - the following should not appear typeset:
+ * \\(x = \\frac{-b\\pm\\sqrt{b^2-4ac}}{2a}\\)
+ */
 export const Math: Story = {
   args: {
     requestOpts: { apiUrl: TEST_API_STREAMING },
@@ -139,13 +145,97 @@ $$
 x = \\frac{-b\\pm\\sqrt{b^2-4ac}}{2a}
 $$
 
+Math is rendered using MathJax only if the \`useMathJax\` prop is set to true.
+`,
+      },
+    ],
+    useMathJax: true,
+  },
+}
+
+/**
+ * Ensures that LaTeX delimiters `\\(...\\)` and `\\[...\\]`  are rendered correctly.
+ *
+ * better-react-mathjax expects TeX delimiters `$...$` or `$$...$$`, though the LLMs may produce LaTeX delimiters despite instruction.
+ *
+ */
+export const MathLatexDelimiters: Story = {
+  args: {
+    requestOpts: { apiUrl: TEST_API_STREAMING },
+    entryScreenEnabled: false,
+    conversationStarters: [],
+    initialMessages: [
+      {
+        role: "assistant",
+        content: `Some inline math: \\(x = \\frac{-b\\pm\\sqrt{b^2-4ac}}{2a}\\)
+
+And some block math:
+
+\\[
+x = \\frac{-b\\pm\\sqrt{b^2-4ac}}{2a}
+\\]
+
 Math is rendered using MathJax only if the \`useMathJax\` prop is set to true.`,
       },
     ],
     useMathJax: true,
   },
 }
 
+/**
+ * > **⚠ NOTE:** Although the `boldsymbol` package is lazily loaded by default,
+ * > we have found that on slow networks a race condition can occur resulting in
+ * > un-rendered boldsymbol commands. This behavior can be reproduced by removing
+ * > the `loader` configuration that explicitly loads the `boldsymbol` package
+ * > and using a throttled network connection.
+ *
+ * This story demonstrates math rendering of LaTeX commands from specialized
+ * packages. MathJax includes some packages in its base configuration, loads
+ * some packages lazily on demand ("autoloads"), and requires some packages to
+ * be explicitly loaded.
+ *
+ * The responses here demonstrate commands from the following packages:
+ *
+ * | Package    | Included by default? | Lazy loaded by default? |
+ * |------------|----------------------|-------------------------|
+ * | `ams`        | yes                  | n/a                     |
+ * | `boldsymbol` | no                   | yes                     |
+ * | `physics`    | no                   | no                      |
+ *
+ * See [https://docs.mathjax.org/en/latest/input/tex/macros/index.html](https://docs.mathjax.org/en/latest/input/tex/macros/index.html)
+ *
+ * Therefore, we explicitly include boldsymbol in the `loader` configuration
+ * to force it to be loaded up front.
+ *
+ * ```tsx
+ * mathJaxConfig: {
+ *   loader: { load: ["[tex]/physics"] },
+ *   tex: { packages: { "[+]": ["physics"] } },
+ * }
+ * ```
+ *
+ * Note that multiple \<MathJaxContext\> instances on the page share a single state. If they have different configs, only the first is applied.
+ *
+ */
+export const MathWithExtensionPackages: Story = {
+  parameters: {
+    height: "800px",
+  },
+  args: {
+    requestOpts: { apiUrl: TEST_API_STREAMING_MATH },
+    entryScreenEnabled: false,
+    conversationStarters: [],
+    initialMessages: [
+      {
+        role: "assistant",
+        content:
+          "Ask me something and I'll respond with math that includes TeX symbols that are not in the base MathJax package",
+      },
+    ],
+    useMathJax: true,
+  },
+}
+
 export const SimpleOrderedList: Story = {
   args: {
     requestOpts: { apiUrl: TEST_API_STREAMING },