Provide messages on assignment selection as prop

Author: jonkafton

src/components/AiChat/AiChat.stories.tsx

@@ -110,7 +110,7 @@ export const JsonResponses: Story = {
   },
 }
 
-export const Temp: Story = {
+export const AssignmentSelection: Story = {
   args: {
     requestOpts: {
       apiUrl: TEST_API_STREAMING,
@@ -126,30 +126,15 @@ export const Temp: Story = {
         role: "assistant",
       },
     ],
-    conversationStarters: STARTERS,
+    conversationStarters: [],
     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: [
+    problemSetListUrl: TEST_API_PROBLEM_SET_LIST,
+    problemSetInitialMessages: [
       {
-        content:
-          "Hi! Please select an assignment from the dropdown menu to begin.",
         role: "assistant",
+        content: "Which question are you working on?",
       },
     ],
-    conversationStarters: [],
-    entryScreenEnabled: false,
-    problemSetListUrl: TEST_API_PROBLEM_SET_LIST,
   },
 }
 

src/components/AiChat/AiChat.tsx

@@ -77,6 +77,9 @@ const AssignmentSelect = styled(SimpleSelectField)({
   "> div": {
     width: "inherit",
   },
+  label: {
+    display: "none",
+  },
 })
 
 const MessagesContainer = styled(ScrollSnap)<{ externalScroll: boolean }>(
@@ -235,6 +238,7 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
   useMathJax = false,
   onSubmit,
   problemSetListUrl,
+  problemSetInitialMessages,
   ...others // Could contain data attributes
 }) => {
   const containerRef = useRef<HTMLDivElement>(null)
@@ -257,6 +261,7 @@ const AiChatDisplay: FC<AiChatDisplayProps> = ({
     initialMessages,
     status,
     additionalBody,
+    setMessages,
     setAdditionalBody,
   } = useAiChat()
 
@@ -298,6 +303,14 @@ 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 })
   }
 

src/components/AiChat/AiChatContext.tsx

@@ -112,17 +112,6 @@ 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
       /**
@@ -132,9 +121,9 @@ const AiChatProvider: React.FC<AiChatContextProps> = ({
       value={{
         initialMessages,
         messages,
-        additionalBody,
-        setAdditionalBody: _setAdditionalBody,
         setMessages,
+        additionalBody,
+        setAdditionalBody,
         ...others,
       }}
     >

src/components/AiChat/types.ts

@@ -124,9 +124,17 @@ type AiChatDisplayProps = {
 
   /**
    * 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: {