feat: show the section that is streaming in during testing

apps/nextjs/src/components/AppComponents/Chat/Chat/hooks/useAilaStreamingStatus.ts

@@ -1,10 +1,52 @@
 import { useMemo, useEffect } from "react";
 
+import type { LessonPlanKeys } from "@oakai/aila/src/protocol/schema";
+import { LessonPlanKeysSchema } from "@oakai/aila/src/protocol/schema";
 import { aiLogger } from "@oakai/logger";
 import type { Message } from "ai";
 
 const log = aiLogger("chat");
 
+function findStreamingSections(message: Message | undefined): {
+  streamingSections: LessonPlanKeys[];
+  streamingSection: LessonPlanKeys | undefined;
+  content: string | undefined;
+} {
+  if (!message?.content) {
+    return {
+      streamingSections: [],
+      streamingSection: undefined,
+      content: undefined,
+    };
+  }
+  const { content } = message;
+  const pathMatches: RegExpExecArray[] = [];
+  let match: RegExpExecArray | null;
+  const regex = /"path":"\/([^/"]*)(?:\/|")"/g;
+  let startIndex = 0;
+  while ((match = regex.exec(content.slice(startIndex))) !== null) {
+    pathMatches.push(match);
+    startIndex += match.index + match[0].length;
+    if (pathMatches.length > 100) {
+      log.warn("Too many path matches found, stopping search");
+      break;
+    }
+  }
+
+  const streamingSections: LessonPlanKeys[] = pathMatches
+    .map((match) => match[1])
+    .filter((i): i is string => typeof i === "string")
+    .map((section) => {
+      const result = LessonPlanKeysSchema.safeParse(section);
+      return result.success ? result.data : undefined;
+    })
+    .filter((section): section is LessonPlanKeys => section !== undefined);
+  const streamingSection: LessonPlanKeys | undefined =
+    streamingSections[streamingSections.length - 1];
+
+  return { streamingSections, streamingSection, content };
+}
+
 export type AilaStreamingStatus =
   | "Loading"
   | "RequestMade"
@@ -18,36 +60,49 @@ export const useAilaStreamingStatus = ({
 }: {
   isLoading: boolean;
   messages: Message[];
-}): AilaStreamingStatus => {
-  const ailaStreamingStatus = useMemo<AilaStreamingStatus>(() => {
-    const moderationStart = "MODERATION_START";
-    const chatStart = "CHAT_START";
-    if (messages.length === 0) return "Idle";
+}): {
+  status: AilaStreamingStatus;
+  streamingSection: LessonPlanKeys | undefined;
+  streamingSections: LessonPlanKeys[] | undefined;
+} => {
+  const { status, streamingSection, streamingSections } = useMemo(() => {
+    const moderationStart = `MODERATION_START`;
+    const chatStart = `CHAT_START`;
+    if (messages.length === 0)
+      return {
+        status: "Idle" as AilaStreamingStatus,
+        streamingSection: undefined,
+      };
     const lastMessage = messages[messages.length - 1];
 
+    let status: AilaStreamingStatus = "Idle";
+    const { streamingSections, streamingSection, content } =
+      findStreamingSections(lastMessage);
+
     if (isLoading) {
-      if (!lastMessage) return "Loading";
-      const { content } = lastMessage;
-      if (lastMessage.role === "user") {
-        return "RequestMade";
-      } else if (content.includes(moderationStart)) {
-        return "Moderating";
-      } else if (
-        content.includes('"type":"prompt"') ||
-        content.includes('\\"type\\":\\"prompt\\"')
-      ) {
-        return "StreamingChatResponse";
-      } else if (content.includes(chatStart)) {
-        return "StreamingLessonPlan";
+      if (!lastMessage || !content) {
+        status = "Loading";
+      } else {
+        if (lastMessage.role === "user") {
+          status = "RequestMade";
+        } else if (content.includes(moderationStart)) {
+          status = "Moderating";
+        } else if (content.includes(`"type":"text"`)) {
+          status = "StreamingChatResponse";
+        } else if (content.includes(chatStart)) {
+          status = "StreamingLessonPlan";
+        } else {
+          status = "Loading";
+        }
       }
-      return "Loading";
     }
-    return "Idle";
+
+    return { status, streamingSections, streamingSection };
   }, [isLoading, messages]);
 
   useEffect(() => {
-    log.info("ailaStreamingStatus set:", ailaStreamingStatus);
-  }, [ailaStreamingStatus]);
+    log.info("ailaStreamingStatus set:", status);
+  }, [status]);
 
-  return ailaStreamingStatus;
+  return { status, streamingSection, streamingSections };
 };

apps/nextjs/src/components/AppComponents/Chat/chat-lhs-header.tsx

@@ -28,6 +28,9 @@ const ChatLhsHeader = ({
             <div data-testid="chat-aila-streaming-status">
               {chat.ailaStreamingStatus}
             </div>
+            <div data-testid="chat-aila-streaming-section">
+              {chat.streamingSection}
+            </div>
           </div>
         )}
         <ChatButton

apps/nextjs/src/components/AppComponents/Chat/chat-list.tsx

@@ -4,6 +4,7 @@ import type { Dispatch, SetStateAction } from "react";
 import { useCallback, useEffect, useRef, useState } from "react";
 
 import type { PersistedModerationBase } from "@oakai/core/src/utils/ailaModeration/moderationSchema";
+import { camelCaseToTitleCase } from "@oakai/core/src/utils/camelCaseConversion";
 import { OakBox, OakFlex, OakIcon, OakSpan } from "@oaknational/oak-components";
 import type { Message } from "ai";
 import Link from "next/link";
@@ -126,13 +127,17 @@ export const ChatMessagesDisplay = ({
   ailaStreamingStatus: AilaStreamingStatus;
   demo: DemoContextProps;
 }) => {
-  const { lessonPlan, isStreaming } = useLessonChat();
+  const { lessonPlan, isStreaming, streamingSection } = useLessonChat();
   const { setDialogWindow } = useDialog();
   const { totalSections, totalSectionsComplete } = useProgressForDownloads({
     lessonPlan,
     isStreaming,
   });
 
+  const workingOnItMessage = streamingSection
+    ? `Editing ${camelCaseToTitleCase(streamingSection)}…`
+    : "Working on it…";
+
   return (
     <>
       {messages.map((message) => {
@@ -159,7 +164,7 @@ export const ChatMessagesDisplay = ({
                 message={{
                   id: "working-on-it-initial",
                   role: "assistant",
-                  content: "Working on it…",
+                  content: workingOnItMessage,
                 }}
                 lastModeration={lastModeration}
                 persistedModerations={[]}
@@ -186,7 +191,7 @@ export const ChatMessagesDisplay = ({
                   ? {
                       id: "working-on-it-initial",
                       role: "assistant",
-                      content: "Working on it…",
+                      content: workingOnItMessage,
                     }
                   : message
               }
@@ -209,7 +214,7 @@ export const ChatMessagesDisplay = ({
               message={{
                 id: "working-on-it-initial",
                 role: "assistant",
-                content: "Working on it…",
+                content: workingOnItMessage,
               }}
               lastModeration={lastModeration}
               persistedModerations={[]}

apps/nextjs/src/components/ContextProviders/ChatProvider.tsx

@@ -13,6 +13,7 @@ import { generateMessageId } from "@oakai/aila/src/helpers/chat/generateMessageI
 import { parseMessageParts } from "@oakai/aila/src/protocol/jsonPatchProtocol";
 import type {
   AilaPersistedChat,
+  LessonPlanKeys,
   LooseLessonPlan,
 } from "@oakai/aila/src/protocol/schema";
 import { isToxic } from "@oakai/core/src/utils/ailaModeration/helpers";
@@ -64,6 +65,8 @@ export type ChatContextProps = {
   queuedUserAction: string | null;
   queueUserAction: (action: string) => void;
   executeQueuedAction: () => Promise<void>;
+  streamingSection: LessonPlanKeys | undefined;
+  streamingSections: LessonPlanKeys[] | undefined;
 };
 
 const ChatContext = createContext<ChatContextProps | null>(null);
@@ -382,7 +385,11 @@ export function ChatProvider({ id, children }: Readonly<ChatProviderProps>) {
       ? lastModeration
       : toxicInitialModeration;
 
-  const ailaStreamingStatus = useAilaStreamingStatus({ isLoading, messages });
+  const {
+    status: ailaStreamingStatus,
+    streamingSection,
+    streamingSections,
+  } = useAilaStreamingStatus({ isLoading, messages });
 
   useEffect(() => {
     if (toxicModeration) {
@@ -416,6 +423,8 @@ export function ChatProvider({ id, children }: Readonly<ChatProviderProps>) {
       queuedUserAction,
       queueUserAction,
       executeQueuedAction,
+      streamingSection,
+      streamingSections,
     }),
     [
       id,
@@ -441,6 +450,8 @@ export function ChatProvider({ id, children }: Readonly<ChatProviderProps>) {
       queuedUserAction,
       queueUserAction,
       executeQueuedAction,
+      streamingSection,
+      streamingSections,
     ],
   );