Assignment list / problem set control

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

@@ -123,6 +123,39 @@ const meta: Meta<typeof AiDrawerManager> = {
           },
         }}
       />
+      <h3>Tutor Bot with assignment selection</h3>
+      <IFrame
+        payload={{
+          blockType: "problem",
+          blockUsageKey: "problem-frame-assignment",
+          trackingUrl: TRACKING_EVENTS_ENDPOINT,
+          chat: {
+            apiUrl: TEST_API_STREAMING,
+            initialMessages: INITIAL_MESSAGES,
+            conversationStarters: STARTERS,
+            entryScreenEnabled: false,
+            entryScreenTitle: "AskTIM about this problem",
+            assignmentList: [
+              {
+                displayName: "Assignment 1",
+                usageKey: "assignment-1",
+              },
+              {
+                displayName: "Assignment 2",
+                usageKey: "assignment-2",
+              },
+              {
+                displayName: "Assignment 3",
+                usageKey: "assignment-3",
+              },
+              {
+                displayName: "Assignment 4",
+                usageKey: "assignment-4",
+              },
+            ],
+          },
+        }}
+      />
       <h3>Video Drawer</h3>
       <p>
         The chat entry screen is shown by default for the video blocks Tutor
@@ -185,7 +218,6 @@ export const AiDrawerManagerStory: Story = {
         }),
         http.post(TRACKING_EVENTS_ENDPOINT, async ({ request }) => {
           const body = await request.json()
-          console.log("TrackingEvent", body)
           return HttpResponse.json({ success: true })
         }),
         ...handlers,

src/bundles/AiDrawer/AiDrawerManager.test.tsx

@@ -67,7 +67,6 @@ const trackingEvent = jest.fn()
 const assertTrackingEvent = (...data: unknown[]) => {
   expect(trackingEvent).toHaveBeenCalledTimes(data.length)
   data.forEach((eventData) => {
-    console.log(eventData)
     expect(trackingEvent).toHaveBeenCalledWith(eventData)
   })
   trackingEvent.mockClear()

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,49 @@ export const JsonResponses: Story = {
   },
 }
 
+export const Temp: 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: STARTERS,
+    entryScreenEnabled: false,
+  },
+}
+
+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,
+  },
+}
+
 const ScrollComponent: FC<AiChatProps> = (args) => {
   const ref = useRef<HTMLDivElement>(null)
   const [scrollElement, setScrollElement] = useState<HTMLDivElement | null>(

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,13 @@ const ChatContainer = styled.div<{ externalScroll: boolean }>(
   }),
 )
 
+const AssignmentSelect = styled(SimpleSelectField)(({ theme }) => ({
+  width: "295px",
+  "> div": {
+    width: "inherit",
+  },
+}))
+
 const MessagesContainer = styled(ScrollSnap)<{ externalScroll: boolean }>(
   ({ externalScroll }) => ({
     display: "flex",
@@ -224,12 +234,16 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
   ref,
   useMathJax = false,
   onSubmit,
+  problemSetListUrl,
   ...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 +256,8 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
     error,
     initialMessages,
     status,
+    additionalBody,
+    setAdditionalBody,
   } = useAiChat()
 
   useScrollSnap({
@@ -281,6 +297,10 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
     })
   }
 
+  const onProblemSetChange = (event: SelectChangeEvent<string | string[]>) => {
+    setAdditionalBody?.({ problem_set_title: event.target.value as string })
+  }
+
   const lastMsg = messages[messages.length - 1]
 
   const externalScroll = !!scrollElement
@@ -317,7 +337,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,
@@ -94,13 +112,31 @@ const AiChatProvider: React.FC<AiChatContextProps> = ({
     })
   }, [parseContent, unparsed, initialMessages])
 
+  const _setAdditionalBody = (body: Record<string, string>) => {
+    setMessages([
+      {
+        role: "assistant",
+        content: "Which question are you working on?",
+        id: "initial-0",
+      },
+    ])
+    setAdditionalBody(body)
+  }
+
   return (
     <AiChatContext.Provider
       /**
        * Ensure that child state is reset when chatId changes.
        */
       key={chatId}
-      value={{ initialMessages, messages, ...others }}
+      value={{
+        initialMessages,
+        messages,
+        additionalBody,
+        setAdditionalBody: _setAdditionalBody,
+        setMessages,
+        ...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>
   )
 }