refactor AiChat to Provider + Display

Author: ChristopherChudzicki

src/ai.ts

@@ -1,3 +1,8 @@
-export { AiChat } from "./components/AiChat/AiChat"
-export type { AiChatProps } from "./components/AiChat/AiChat"
-export type { AiChatMessage } from "./components/AiChat/types"
+export { AiChat, AiChatDisplay } from "./components/AiChat/AiChat"
+export { AiChatProvider, useAiChat } from "./components/AiChat/AiChatContext"
+export type {
+  AiChatMessage,
+  AiChatContextProps,
+  AiChatDisplayProps,
+  AiChatProps,
+} from "./components/AiChat/types"

src/bundles/RemoteTutorDrawer/RemoteTutorDrawer.stories.tsx

@@ -74,6 +74,7 @@ const IFrame = ({ payload }: { payload: InitPayload }) => {
 
 const meta: Meta<typeof RemoteTutorDrawer> = {
   title: "smoot-design/AI/RemoteTutorDrawer",
+  component: RemoteTutorDrawer,
   render: ({ target }, { parameters: { payload } }) => (
     <>
       <IFrame payload={payload} />

src/bundles/RemoteTutorDrawer/RemoteTutorDrawer.tsx

@@ -14,8 +14,7 @@ import Typography from "@mui/material/Typography"
 import TabContext from "@mui/lab/TabContext"
 import TabPanel from "@mui/lab/TabPanel"
 import { AiChat } from "../../components/AiChat/AiChat"
-import { AiChatMessage } from "../../components/AiChat/types"
-import type { AiChatProps } from "../../components/AiChat/AiChat"
+import type { AiChatProps, AiChatMessage } from "../../components/AiChat/types"
 import { ActionButton } from "../../components/Button/ActionButton"
 import { FlashcardsScreen } from "./FlashcardsScreen"
 import type { Flashcard } from "./FlashcardsScreen"

src/components/AiChat/AiChat.tsx

@@ -1,22 +1,21 @@
 import * as React from "react"
-import { useEffect, useRef, useState, useMemo } from "react"
+import { useEffect, useRef, useState } from "react"
 import type { FC } from "react"
 import styled from "@emotion/styled"
 import Typography from "@mui/material/Typography"
 import classNames from "classnames"
 import { RiSendPlaneFill, RiStopFill, RiMoreFill } from "@remixicon/react"
 import { Input, AdornmentButton } from "../Input/Input"
-import type { AiChatMessage, AiChatProps } from "./types"
-import { EntryScreen } from "./EntryScreen"
+import type { AiChatDisplayProps, AiChatProps } from "./types"
 import Markdown from "react-markdown"
 import { ScrollSnap } from "../ScrollSnap/ScrollSnap"
 import { SrAnnouncer } from "../SrAnnouncer/SrAnnouncer"
 import { VisuallyHidden } from "../VisuallyHidden/VisuallyHidden"
 import { Alert } from "../Alert/Alert"
 import { ChatTitle } from "./ChatTitle"
-import { useAiChat } from "./utils"
+import { AiChatProvider, useAiChat } from "./AiChatContext"
 import { useScrollSnap } from "../ScrollSnap/useScrollSnap"
-import type { Message } from "@ai-sdk/react"
+import { EntryScreen } from "./EntryScreen"
 
 const classes = {
   root: "MitAiChat--root",
@@ -179,78 +178,48 @@ const Disclaimer = styled(Typography)(({ theme }) => ({
   textAlign: "center",
 }))
 
-const AiChat: FC<AiChatProps> = ({
-  entryScreenTitle,
-  entryScreenEnabled = true,
+const AiChatDisplay: FC<AiChatDisplayProps> = ({
   conversationStarters,
-  initialMessages: _initialMessages,
   askTimTitle,
-  requestOpts,
-  parseContent,
+  entryScreenEnabled = true,
+  entryScreenTitle,
   srLoadingMessages,
   placeholder = "",
   className,
   scrollElement,
-  chatId,
   ref,
   ...others // Could contain data attributes
 }) => {
   const containerRef = useRef<HTMLDivElement>(null)
   const messagesContainerRef = useRef<HTMLDivElement>(null)
-  const [showEntryScreen, setShowEntryScreen] = useState(entryScreenEnabled)
   const chatScreenRef = useRef<HTMLDivElement>(null)
-  const [initialMessages, setInitialMessages] = useState<AiChatMessage[]>()
   const promptInputRef = useRef<HTMLDivElement>(null)
 
   const {
-    messages: unparsed,
+    messages,
     input,
     handleInputChange,
     handleSubmit,
     append,
     isLoading,
     stop,
     error,
-  } = useAiChat(requestOpts, {
     initialMessages,
-    id: chatId,
-  })
+  } = useAiChat()
 
   useScrollSnap({
     scrollElement: scrollElement || messagesContainerRef.current,
     contentElement: scrollElement ? messagesContainerRef.current : null,
     threshold: 200,
   })
 
-  useEffect(() => {
-    if (_initialMessages) {
-      const prefix = Math.random().toString().slice(2)
-      setInitialMessages(
-        _initialMessages.map((m, i) => ({
-          ...m,
-          id: `initial-${prefix}-${i}`,
-        })),
-      )
-    }
-  }, [_initialMessages])
-
+  const [showEntryScreen, setShowEntryScreen] = useState(entryScreenEnabled)
   useEffect(() => {
     if (!showEntryScreen) {
       promptInputRef.current?.querySelector("input")?.focus()
     }
   }, [showEntryScreen])
 
-  const messages = useMemo(() => {
-    const initial = initialMessages?.map((m) => m.id)
-    return unparsed.map((m: Message) => {
-      if (m.role === "assistant" && !initial?.includes(m.id)) {
-        const content = parseContent ? parseContent(m.content) : m.content
-        return { ...m, content }
-      }
-      return m
-    })
-  }, [parseContent, unparsed, initialMessages])
-
   const showStarters = messages.length === (initialMessages?.length || 0)
 
   const waiting =
@@ -270,19 +239,7 @@ const AiChat: FC<AiChatProps> = ({
   const externalScroll = !!scrollElement
 
   return (
-    <Container
-      className={className}
-      ref={containerRef}
-      /**
-       * Changing the `useChat` chatId seems to persist some state between
-       * hook calls. This can cause strange effects like loading API responses
-       * for previous chatId into new chatId.
-       *
-       * To avoid this, let's change the key, this will force React to make a new component
-       * not sharing any of the old state.
-       */
-      key={chatId}
-    >
+    <Container className={className} ref={containerRef}>
       {showEntryScreen ? (
         <EntryScreen
           className={classes.entryScreenContainer}
@@ -318,7 +275,7 @@ const AiChat: FC<AiChatProps> = ({
               externalScroll={externalScroll}
               ref={messagesContainerRef}
             >
-              {messages.map((m: Message) => (
+              {messages.map((m) => (
                 <MessageRow
                   key={m.id}
                   data-chat-role={m.role}
@@ -442,5 +399,23 @@ const AiChat: FC<AiChatProps> = ({
   )
 }
 
-export { AiChat }
-export type { AiChatProps }
+const AiChat: FC<AiChatProps> = ({
+  requestOpts,
+  initialMessages,
+  chatId,
+  parseContent,
+  ...displayProps
+}) => {
+  return (
+    <AiChatProvider
+      requestOpts={requestOpts}
+      chatId={chatId}
+      initialMessages={initialMessages}
+      parseContent={parseContent}
+    >
+      <AiChatDisplay {...displayProps} />
+    </AiChatProvider>
+  )
+}
+
+export { AiChatDisplay, AiChat }

src/components/AiChat/AiChatContext.stories.tsx

@@ -0,0 +1,96 @@
+import * as React from "react"
+import type { Meta, StoryObj } from "@storybook/react"
+import { AiChatDisplay } from "./AiChat"
+import { AiChatProvider, useAiChat } from "./AiChatContext"
+import type { AiChatProps } from "./types"
+import styled from "@emotion/styled"
+import { handlers } from "./test-utils/api"
+import Typography from "@mui/material/Typography"
+
+const TEST_API_STREAMING = "http://localhost:4567/streaming"
+
+const INITIAL_MESSAGES: AiChatProps["initialMessages"] = [
+  {
+    content: "Hi! What are you interested in learning about?",
+    role: "assistant",
+  },
+]
+
+const STARTERS = [
+  { content: "I'm interested in quantum computing" },
+  { content: "I want to understand global warming. " },
+  { content: "I am curious about AI applications for business" },
+]
+
+const Container = styled.div({
+  width: "100%",
+  height: "400px",
+  position: "relative",
+})
+
+const MessageCounter = () => {
+  const { messages } = useAiChat()
+
+  return (
+    <Typography variant="subtitle1">
+      Message count: {messages.length} (Provided by <code>AiChatContext</code>)
+    </Typography>
+  )
+}
+
+/**
+ * AiChatProvider provides state and functions for managing chat. The higher-level
+ * `AiChat` component is a wrapper around this provider and the `AiChatDisplay`,
+ *  roughly.
+ *
+ * If you need to access chat state outside of the chat display, you can use
+ * `AiChatProvider` directly.
+ */
+const meta: Meta<typeof AiChatProvider> = {
+  title: "smoot-design/AI/AiChatContext",
+  component: AiChatProvider,
+  parameters: {
+    msw: { handlers },
+  },
+  render: (args) => {
+    return (
+      <AiChatProvider {...args}>
+        <MessageCounter />
+        <Container>
+          <AiChatDisplay
+            entryScreenEnabled={false}
+            conversationStarters={STARTERS}
+            placeholder="Type your message here"
+            askTimTitle="Ask TIM"
+          />
+        </Container>
+      </AiChatProvider>
+    )
+  },
+  decorators: (Story) => {
+    return (
+      <Container>
+        <Story />
+      </Container>
+    )
+  },
+  args: {
+    requestOpts: { apiUrl: TEST_API_STREAMING },
+    initialMessages: INITIAL_MESSAGES,
+  },
+  argTypes: {
+    initialMessages: {
+      control: { type: "object", disable: true },
+    },
+    requestOpts: {
+      control: { type: "object", disable: true },
+      table: { readonly: true }, // See above
+    },
+  },
+}
+
+export default meta
+
+type Story = StoryObj<typeof AiChatProvider>
+
+export const StreamingResponses: Story = {}