Skip to content

Commit 7a82853

Browse files
jonkaftonjonkafton
committed
Move MathJax context into the AiChat component so it's always available
1 parent c235ad8 commit 7a82853

File tree

7 files changed

+101
-103
lines changed

7 files changed

+101
-103
lines changed

‎src/bundles/AiDrawer/AiDrawer.stories.tsx

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import type { Meta, StoryObj } from "@storybook/nextjs"
44
import { http, HttpResponse } from "msw"
55
import { handlers } from "../../components/AiChat/test-utils/api"
66
import { AiDrawer, AiDrawerSettings } from "./AiDrawer"
7-
import { MathJaxContext } from "better-react-mathjax"
87
import Button from "@mui/material/Button"
98

109
const TEST_API_STREAMING = "http://localhost:4567/streaming"
@@ -36,13 +35,11 @@ const meta: Meta<typeof AiDrawer> = {
3635
</Button>
3736
<p>Message data:</p>
3837
<pre>{JSON.stringify({ settings }, null, 2)}</pre>
39-
<MathJaxContext>
40-
<AiDrawer
41-
settings={settings}
42-
open={open}
43-
onClose={() => setOpen(false)}
44-
/>
45-
</MathJaxContext>
38+
<AiDrawer
39+
settings={settings}
40+
open={open}
41+
onClose={() => setOpen(false)}
42+
/>
4643
</>
4744
)
4845
},

‎src/bundles/AiDrawer/AiDrawerManager.stories.tsx

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import { http, HttpResponse } from "msw"
66
import { handlers } from "../../components/AiChat/test-utils/api"
77
import { AiDrawerManager } from "./AiDrawerManager"
88
import type { AiDrawerInitMessage } from "./AiDrawerManager"
9-
import { MathJaxContext } from "better-react-mathjax"
109

1110
type InitPayload = AiDrawerInitMessage["payload"]
1211

@@ -162,9 +161,7 @@ const meta: Meta<typeof AiDrawerManager> = {
162161
},
163162
}}
164163
/>
165-
<MathJaxContext>
166-
<AiDrawerManager messageOrigin="http://localhost:6006" />
167-
</MathJaxContext>
164+
<AiDrawerManager messageOrigin="http://localhost:6006" />
168165
</div>
169166
),
170167
}

‎src/bundles/AiDrawer/AiDrawerManager.tsx

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ import * as React from "react"
22
import { useEffect, useState } from "react"
33
import { AiDrawer } from "./AiDrawer"
44
import type { AiDrawerProps, AiDrawerSettings } from "./AiDrawer"
5-
import { MathJaxContext } from "better-react-mathjax"
65

76
type AiDrawerInitMessage = {
87
type: "smoot-design::ai-drawer-open" | "smoot-design::tutor-drawer-open" // ("smoot-design::tutor-drawer-open" is legacy)
@@ -109,7 +108,7 @@ const AiDrawerManager = ({
109108
}
110109

111110
return (
112-
<MathJaxContext>
111+
<>
113112
{Object.values(drawerStates).map(({ key, open, payload }) => {
114113
const { trackingUrl, ...settings } = payload
115114
return (
@@ -147,7 +146,7 @@ const AiDrawerManager = ({
147146
/>
148147
)
149148
})}
150-
</MathJaxContext>
149+
</>
151150
)
152151
}
153152

‎src/bundles/aiDrawerManager.tsx

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import {
88
} from "../components/ThemeProvider/ThemeProvider"
99
import { CacheProvider } from "@emotion/react"
1010
import createCache from "@emotion/cache"
11-
import { MathJaxContext } from "better-react-mathjax"
1211

1312
/**
1413
* Renders the AiDrawerManager to the page.
@@ -33,9 +32,7 @@ const init = (opts: AiDrawerManagerProps) => {
3332
createRoot(container).render(
3433
<CacheProvider value={cache}>
3534
<ThemeProvider theme={theme}>
36-
<MathJaxContext>
37-
<AiDrawerManager {...opts} />
38-
</MathJaxContext>
35+
<AiDrawerManager {...opts} />
3936
</ThemeProvider>
4037
</CacheProvider>,
4138
)

‎src/components/AiChat/AiChat.stories.tsx

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import type { AiChatProps } from "./types"
66
import styled from "@emotion/styled"
77
import { handlers } from "./test-utils/api"
88
import { FC, useEffect, useRef, useState } from "react"
9-
import { MathJaxContext } from "better-react-mathjax"
109

1110
const TEST_API_STREAMING = "http://localhost:4567/streaming"
1211
const TEST_API_JSON = "http://localhost:4567/json"
@@ -62,11 +61,9 @@ const meta: Meta<typeof AiChat> = {
6261
render: (args) => <AiChat {...args} />,
6362
decorators: (Story, context) => {
6463
return (
65-
<MathJaxContext>
66-
<Container>
67-
<Story key={String(context.args.entryScreenEnabled)} />
68-
</Container>
69-
</MathJaxContext>
64+
<Container>
65+
<Story key={String(context.args.entryScreenEnabled)} />
66+
</Container>
7067
)
7168
},
7269
args: {

‎src/components/AiChat/AiChat.tsx

Lines changed: 86 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import * as React from "react"
2-
import { useEffect, useRef, useState } from "react"
2+
import { useEffect, useRef, useState, Suspense } from "react"
33
import type { FC } from "react"
44
import styled from "@emotion/styled"
55
import Typography from "@mui/material/Typography"
@@ -21,6 +21,17 @@ import EllipsisIcon from "./EllipsisIcon"
2121
import { SimpleSelectField } from "../SimpleSelect/SimpleSelect"
2222
import { useFetch } from "./utils"
2323
import { SelectChangeEvent } from "@mui/material/Select"
24+
import { MathJaxContext } from "better-react-mathjax"
25+
26+
const ConditionalMathJaxWrapper: React.FC<{
27+
useMathJax: boolean
28+
children: React.ReactNode
29+
}> = ({ useMathJax, children }) => {
30+
if (!useMathJax) {
31+
return <>{children}</>
32+
}
33+
return <MathJaxContext>{children}</MathJaxContext>
34+
}
2435

2536
const classes = {
2637
root: "MitAiChat--root",
@@ -377,81 +388,84 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
377388
) : null
378389
}
379390
/>
380-
381-
<MessagesContainer
382-
className={classes.messagesContainer}
383-
externalScroll={externalScroll}
384-
ref={messagesContainerRef}
385-
>
386-
{messages.map((m: Message, i) => {
387-
// Our Markdown+Mathjax has issues when rendering streaming display math
388-
// Force a re-render of the last (streaming) message when it's done loading.
389-
const key =
390-
i === messages.length - 1 && isLoading
391-
? `isLoading-${m.id}`
392-
: m.id
393-
return (
391+
<ConditionalMathJaxWrapper useMathJax={useMathJax}>
392+
<MessagesContainer
393+
className={classes.messagesContainer}
394+
externalScroll={externalScroll}
395+
ref={messagesContainerRef}
396+
>
397+
{messages.map((m: Message, i) => {
398+
// Our Markdown+Mathjax has issues when rendering streaming display math
399+
// Force a re-render of the last (streaming) message when it's done loading.
400+
const key =
401+
i === messages.length - 1 && isLoading
402+
? `isLoading-${m.id}`
403+
: m.id
404+
return (
405+
<MessageRow
406+
key={key}
407+
data-chat-role={m.role}
408+
className={classNames(classes.messageRow, {
409+
[classes.messageRowUser]: m.role === "user",
410+
[classes.messageRowAssistant]: m.role === "assistant",
411+
})}
412+
>
413+
<Message className={classes.message}>
414+
<VisuallyHidden as={m.role === "user" ? "h5" : "h6"}>
415+
{m.role === "user"
416+
? "You said: "
417+
: "Assistant said: "}
418+
</VisuallyHidden>
419+
{useMathJax ? (
420+
<Markdown enableMathjax={true}>
421+
{replaceMathjax(m.content)}
422+
</Markdown>
423+
) : (
424+
<Markdown>{m.content}</Markdown>
425+
)}
426+
</Message>
427+
</MessageRow>
428+
)
429+
})}
430+
{showStarters ? (
431+
<StarterContainer>
432+
{conversationStarters?.map((m) => (
433+
<Starter
434+
className={classes.conversationStarter}
435+
key={m.content}
436+
onClick={() => {
437+
scrollToBottom()
438+
append({ role: "user", content: m.content })
439+
onSubmit?.(m.content, {
440+
source: "conversation-starter",
441+
})
442+
}}
443+
>
444+
{m.content}
445+
</Starter>
446+
))}
447+
</StarterContainer>
448+
) : null}
449+
{waiting ? (
394450
<MessageRow
395-
key={key}
396-
data-chat-role={m.role}
397-
className={classNames(classes.messageRow, {
398-
[classes.messageRowUser]: m.role === "user",
399-
[classes.messageRowAssistant]: m.role === "assistant",
400-
})}
451+
className={classNames(
452+
classes.messageRow,
453+
classes.messageRowAssistant,
454+
)}
455+
key={"loading"}
401456
>
402-
<Message className={classes.message}>
403-
<VisuallyHidden as={m.role === "user" ? "h5" : "h6"}>
404-
{m.role === "user" ? "You said: " : "Assistant said: "}
405-
</VisuallyHidden>
406-
{useMathJax ? (
407-
<Markdown enableMathjax={true}>
408-
{replaceMathjax(m.content)}
409-
</Markdown>
410-
) : (
411-
<Markdown>{m.content}</Markdown>
412-
)}
457+
<Message>
458+
<StyledEllipsisIcon />
413459
</Message>
414460
</MessageRow>
415-
)
416-
})}
417-
{showStarters ? (
418-
<StarterContainer>
419-
{conversationStarters?.map((m) => (
420-
<Starter
421-
className={classes.conversationStarter}
422-
key={m.content}
423-
onClick={() => {
424-
scrollToBottom()
425-
append({ role: "user", content: m.content })
426-
onSubmit?.(m.content, {
427-
source: "conversation-starter",
428-
})
429-
}}
430-
>
431-
{m.content}
432-
</Starter>
433-
))}
434-
</StarterContainer>
435-
) : null}
436-
{waiting ? (
437-
<MessageRow
438-
className={classNames(
439-
classes.messageRow,
440-
classes.messageRowAssistant,
441-
)}
442-
key={"loading"}
443-
>
444-
<Message>
445-
<StyledEllipsisIcon />
446-
</Message>
447-
</MessageRow>
448-
) : null}
449-
{error ? (
450-
<Alert severity="error" closable>
451-
An unexpected error has occurred.
452-
</Alert>
453-
) : null}
454-
</MessagesContainer>
461+
) : null}
462+
{error ? (
463+
<Alert severity="error" closable>
464+
An unexpected error has occurred.
465+
</Alert>
466+
) : null}
467+
</MessagesContainer>
468+
</ConditionalMathJaxWrapper>
455469
<BottomSection
456470
externalScroll={externalScroll}
457471
className={classes.bottomSection}

‎src/components/AiChat/AiChatMarkdown.stories.tsx

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ import type { Meta, StoryObj } from "@storybook/nextjs"
33
import { AiChat } from "./AiChat"
44
import styled from "@emotion/styled"
55
import { handlers } from "./test-utils/api"
6-
import { MathJaxContext } from "better-react-mathjax"
76

87
const TEST_API_STREAMING = "http://localhost:4567/streaming"
98

@@ -21,11 +20,9 @@ const meta: Meta<typeof AiChat> = {
2120
render: (args) => <AiChat {...args} />,
2221
decorators: (Story, context) => {
2322
return (
24-
<MathJaxContext>
25-
<Container>
26-
<Story key={String(context.args.entryScreenEnabled)} />
27-
</Container>
28-
</MathJaxContext>
23+
<Container>
24+
<Story key={String(context.args.entryScreenEnabled)} />
25+
</Container>
2926
)
3027
},
3128
args: {

0 commit comments

Comments
 (0)