feat: Assignment selection for chat (#155)

Co-authored-by: Asad Ali <asad.ali@arbisoft.com>

Author: jonkafton

package.json

@@ -92,8 +92,10 @@
     "@storybook/addon-links": "^9.0.13",
     "@storybook/addon-onboarding": "^9.0.13",
     "@storybook/addon-webpack5-compiler-swc": "^3.0.0",
+    "@storybook/blocks": "^8.6.14",
     "@storybook/nextjs": "^9.0.13",
     "@storybook/react-webpack5": "^9.0.13",
+    "@storybook/test": "^8.6.14",
     "@swc/jest": "^0.2.37",
     "@testing-library/dom": "^10.4.0",
     "@testing-library/jest-dom": "^6.6.3",

src/bundles/AiDrawer/AiDrawerManager.stories.tsx

@@ -183,9 +183,7 @@ export const AiDrawerManagerStory: Story = {
         http.get(CONTENT_FILE_URL, () => {
           return HttpResponse.json(sampleResponse)
         }),
-        http.post(TRACKING_EVENTS_ENDPOINT, async ({ request }) => {
-          const body = await request.json()
-          console.log("TrackingEvent", body)
+        http.post(TRACKING_EVENTS_ENDPOINT, () => {
           return HttpResponse.json({ success: true })
         }),
         ...handlers,

src/components/AiChat/AiChat.stories.tsx

@@ -1,5 +1,6 @@
 import * as React from "react"
 import type { Meta, StoryObj } from "@storybook/nextjs"
+import { http, HttpResponse } from "msw"
 import { AiChat } from "./AiChat"
 import type { AiChatProps } from "./types"
 import styled from "@emotion/styled"
@@ -9,6 +10,7 @@ import { MathJaxContext } from "better-react-mathjax"
 
 const TEST_API_STREAMING = "http://localhost:4567/streaming"
 const TEST_API_JSON = "http://localhost:4567/json"
+const TEST_API_PROBLEM_SET_LIST = "http://localhost:4567/problem_set_list"
 
 const INITIAL_MESSAGES: AiChatProps["initialMessages"] = [
   {
@@ -41,7 +43,21 @@ const meta: Meta<typeof AiChat> = {
   title: "smoot-design/AI/AiChat",
   component: AiChat,
   parameters: {
-    msw: { handlers },
+    msw: {
+      handlers: [
+        http.get(TEST_API_PROBLEM_SET_LIST, () => {
+          return HttpResponse.json({
+            problem_set_titles: [
+              "Assignment 1",
+              "Assignment 2",
+              "Assignment 3",
+              "Assignment 4",
+            ],
+          })
+        }),
+        ...handlers,
+      ],
+    },
   },
   render: (args) => <AiChat {...args} />,
   decorators: (Story, context) => {
@@ -94,6 +110,34 @@ export const JsonResponses: Story = {
   },
 }
 
+export const AssignmentSelection: Story = {
+  args: {
+    requestOpts: {
+      apiUrl: TEST_API_STREAMING,
+      transformBody: (messages, body) => ({
+        message: messages[messages.length - 1].content,
+        problem_set_title: body?.problem_set_title,
+      }),
+    },
+    initialMessages: [
+      {
+        content:
+          "Hi! Please select an assignment from the dropdown menu to begin.",
+        role: "assistant",
+      },
+    ],
+    conversationStarters: [],
+    entryScreenEnabled: false,
+    problemSetListUrl: TEST_API_PROBLEM_SET_LIST,
+    problemSetInitialMessages: [
+      {
+        role: "assistant",
+        content: "Which question are you working on?",
+      },
+    ],
+  },
+}
+
 const ScrollComponent: FC<AiChatProps> = (args) => {
   const ref = useRef<HTMLDivElement>(null)
   const [scrollElement, setScrollElement] = useState<HTMLDivElement | null>(

src/components/AiChat/AiChat.test.tsx

@@ -179,10 +179,13 @@ describe("AiChat", () => {
     await user.paste("User message")
     await user.click(screen.getByRole("button", { name: "Send" }))
 
-    expect(transformBody).toHaveBeenCalledWith([
-      expect.objectContaining(initialMessages[0]),
-      expect.objectContaining({ content: "User message", role: "user" }),
-    ])
+    expect(transformBody).toHaveBeenCalledWith(
+      [
+        expect.objectContaining(initialMessages[0]),
+        expect.objectContaining({ content: "User message", role: "user" }),
+      ],
+      {},
+    )
     expect(mockFetch).toHaveBeenCalledTimes(1)
     expect(mockFetch).toHaveBeenCalledWith(
       API_URL,

src/components/AiChat/AiChat.tsx

@@ -18,6 +18,9 @@ import { useScrollSnap } from "../ScrollSnap/useScrollSnap"
 import type { Message } from "@ai-sdk/react"
 import Markdown from "./Markdown"
 import EllipsisIcon from "./EllipsisIcon"
+import { SimpleSelectField } from "../SimpleSelect/SimpleSelect"
+import { useFetch } from "./utils"
+import { SelectChangeEvent } from "@mui/material/Select"
 
 const classes = {
   root: "MitAiChat--root",
@@ -69,6 +72,16 @@ const ChatContainer = styled.div<{ externalScroll: boolean }>(
   }),
 )
 
+const AssignmentSelect = styled(SimpleSelectField)({
+  width: "295px",
+  "> div": {
+    width: "inherit",
+  },
+  label: {
+    display: "none",
+  },
+})
+
 const MessagesContainer = styled(ScrollSnap)<{ externalScroll: boolean }>(
   ({ externalScroll }) => ({
     display: "flex",
@@ -224,12 +237,17 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
   ref,
   useMathJax = false,
   onSubmit,
+  problemSetListUrl,
+  problemSetInitialMessages,
   ...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)
+  const { response: problemSetListResponse } = useFetch<{
+    problem_set_titles: string[]
+  }>(problemSetListUrl)
 
   const {
     messages,
@@ -242,6 +260,9 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
     error,
     initialMessages,
     status,
+    additionalBody,
+    setMessages,
+    setAdditionalBody,
   } = useAiChat()
 
   useScrollSnap({
@@ -281,6 +302,18 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
     })
   }
 
+  const onProblemSetChange = (event: SelectChangeEvent<string | string[]>) => {
+    if (problemSetInitialMessages) {
+      setMessages(
+        problemSetInitialMessages.map((message, i) => ({
+          ...message,
+          id: `initial-${i}`,
+        })),
+      )
+    }
+    setAdditionalBody?.({ problem_set_title: event.target.value as string })
+  }
+
   const lastMsg = messages[messages.length - 1]
 
   const externalScroll = !!scrollElement
@@ -317,7 +350,30 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
               askTimTitle={askTimTitle}
               externalScroll={externalScroll}
               className={classNames(className, classes.title)}
+              control={
+                problemSetListResponse?.problem_set_titles?.length ? (
+                  <AssignmentSelect
+                    label="Assignments"
+                    options={[
+                      {
+                        value: "",
+                        label: "Select an assignment",
+                        disabled: true,
+                      },
+                      ...problemSetListResponse.problem_set_titles.map(
+                        (title) => ({
+                          value: title,
+                          label: title,
+                        }),
+                      ),
+                    ]}
+                    value={additionalBody?.problem_set_title ?? ""}
+                    onChange={onProblemSetChange}
+                  />
+                ) : null
+              }
             />
+
             <MessagesContainer
               className={classes.messagesContainer}
               externalScroll={externalScroll}

src/components/AiChat/AiChatContext.tsx

@@ -1,24 +1,29 @@
 import * as React from "react"
 import { useChat, UseChatHelpers } from "@ai-sdk/react"
 import type { RequestOpts, AiChatMessage, AiChatContextProps } from "./types"
-import { useMemo, createContext } from "react"
+import { useMemo, createContext, useState } from "react"
 import retryingFetch from "../../utils/retryingFetch"
 import { getCookie } from "../../utils/getCookie"
 
 const identity = <T,>(x: T): T => x
 
-const getFetcher: (requestOpts: RequestOpts) => typeof fetch =
-  (requestOpts: RequestOpts) => async (url, opts) => {
+const getFetcher: (
+  requestOpts: RequestOpts,
+  additionalBody: Record<string, string>,
+) => typeof fetch =
+  (requestOpts: RequestOpts, additionalBody: Record<string, string> = {}) =>
+  async (url, opts) => {
     if (typeof opts?.body !== "string") {
       console.error("Unexpected body type.")
       return retryingFetch(url, opts)
     }
-    const messages: AiChatMessage[] = JSON.parse(opts?.body).messages
+    const parsedBody = JSON.parse(opts?.body)
+    const messages: AiChatMessage[] = parsedBody.messages
     const transformBody: RequestOpts["transformBody"] =
       requestOpts.transformBody ?? identity
     const options: RequestInit = {
       ...opts,
-      body: JSON.stringify(transformBody(messages)),
+      body: JSON.stringify(transformBody(messages, additionalBody)),
       ...requestOpts.fetchOpts,
       headers: {
         ...opts?.headers,
@@ -43,6 +48,8 @@ const getFetcher: (requestOpts: RequestOpts) => typeof fetch =
  */
 type AiChatContextResult = UseChatHelpers & {
   initialMessages: AiChatMessage[] | null
+  additionalBody?: Record<string, string>
+  setAdditionalBody?: (body: Record<string, string>) => void
 }
 const AiChatContext = createContext<AiChatContextResult | null>(null)
 
@@ -66,8 +73,19 @@ const AiChatProvider: React.FC<AiChatContextProps> = ({
     )
   }, [_initialMessages])
 
-  const fetcher = useMemo(() => getFetcher(requestOpts), [requestOpts])
-  const { messages: unparsed, ...others } = useChat({
+  const [additionalBody, setAdditionalBody] = useState<Record<string, string>>(
+    {},
+  )
+
+  const fetcher = useMemo(
+    () => getFetcher(requestOpts, additionalBody),
+    [requestOpts, additionalBody],
+  )
+  const {
+    messages: unparsed,
+    setMessages,
+    ...others
+  } = useChat({
     api: requestOpts.apiUrl,
     streamProtocol: "text",
     fetch: fetcher,
@@ -100,7 +118,14 @@ const AiChatProvider: React.FC<AiChatContextProps> = ({
        * Ensure that child state is reset when chatId changes.
        */
       key={chatId}
-      value={{ initialMessages, messages, ...others }}
+      value={{
+        initialMessages,
+        messages,
+        setMessages,
+        additionalBody,
+        setAdditionalBody,
+        ...others,
+      }}
     >
       {children}
     </AiChatContext.Provider>

src/components/AiChat/ChatTitle.tsx

@@ -37,17 +37,22 @@ const AskTimTitle = styled.div(({ theme }) => ({
     width: "24px",
     height: "24px",
   },
+  "&& p": {
+    margin: 0,
+  },
 }))
 
 type ChatTitleProps = {
   askTimTitle?: string
   externalScroll?: boolean
+  control?: React.ReactNode
   className?: string
 }
 
 const ChatTitle = ({
   askTimTitle,
   externalScroll,
+  control,
   className,
 }: ChatTitleProps) => {
   if (!askTimTitle) return null
@@ -60,6 +65,7 @@ const ChatTitle = ({
           {askTimTitle}
         </Typography>
       </AskTimTitle>
+      {control}
     </Container>
   )
 }

src/components/AiChat/types.ts

@@ -18,7 +18,10 @@ type RequestOpts = {
    *
    * JSON.stringify is applied to the return value.
    */
-  transformBody?: (messages: AiChatMessage[]) => unknown
+  transformBody?: (
+    messages: AiChatMessage[],
+    body?: Record<string, string>,
+  ) => unknown
   /**
    * Extra options to pass to fetch.
    */
@@ -52,6 +55,9 @@ type AiChatContextProps = {
   initialMessages?: Omit<AiChatMessage, "id">[]
 
   children?: React.ReactNode
+
+  additionalBody?: Record<string, string>
+  setAdditionalBody?: (body: Record<string, string>) => void
 }
 
 type AiChatDisplayProps = {
@@ -106,7 +112,7 @@ type AiChatDisplayProps = {
   scrollElement?: HTMLElement | null
 
   /**
-   * If true, the chat will display math equations using MathJax.
+   * If true, the chat will display math equations using MathJax..
    * Defaults to false.
    */
   useMathJax?: boolean
@@ -116,6 +122,19 @@ type AiChatDisplayProps = {
    */
   autofocus?: boolean
 
+  /**
+   * URL to fetch problem set list for dropdown.
+   *
+   * The problem set selection is passed as the second argument to the `transformBody` function
+   * provided as `{ problem_set_title: string }`.
+   */
+  problemSetListUrl?: string
+
+  /**
+   * Initial messages to display on problem set selection.
+   */
+  problemSetInitialMessages?: Omit<AiChatMessage, "id">[]
+
   onSubmit?: (
     messageText: string,
     meta: {