fix: chat streaming and loading states (WIP)

.vscode/settings.json

@@ -66,6 +66,7 @@
     "hget",
     "hgetall",
     "hmset",
+    "hotfixes",
     "hset",
     "Hubspot",
     "initialisation",
@@ -131,6 +132,7 @@
     "rushstack",
     "Sedar",
     "slidedeck",
+    "sslcert",
     "sslmode",
     "SUBJ",
     "superjson",
@@ -147,6 +149,7 @@
     "trivago",
     "trpc",
     "Turbopack",
+    "TURBOPACK",
     "turborepo",
     "uidotdev",
     "unjudged",

apps/nextjs/src/app/aila/page-contents.tsx

@@ -11,8 +11,8 @@ const ChatPageContents = ({ id }: { id: string }) => {
   return (
     <Layout>
       <LessonPlanTrackingProvider chatId={id}>
-        <ChatProvider key={`chat-${id}`} id={id}>
-          <Chat />
+        <ChatProvider key={`chat-provider-${id}`} id={id}>
+          <Chat key={`chat-${id}`} />
         </ChatProvider>
       </LessonPlanTrackingProvider>
     </Layout>

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

@@ -1,10 +1,42 @@
 import { useMemo, useEffect } from "react";
 
+import type { LessonPlanKeys } from "@oakai/aila/src/protocol/schema";
 import { aiLogger } from "@oakai/logger";
 import type { Message } from "ai";
 
+import { allSectionsInOrder } from "../../../../../lib/lessonPlan/sectionsInOrder";
+
 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,
+    };
+  }
+  log.info("Parsing message content", message.content);
+  const { content } = message;
+  const regex = /"path":"\/([^/"]*)/g;
+  const pathMatches =
+    content
+      .match(regex)
+      ?.map((match) => match.replace(/"path":"\//, "").replace(/"$/, "")) ?? [];
+
+  const streamingSections: LessonPlanKeys[] = pathMatches.filter(
+    (i): i is string =>
+      typeof i === "string" && allSectionsInOrder.includes(i as LessonPlanKeys),
+  ) as LessonPlanKeys[];
+  const streamingSection: LessonPlanKeys | undefined =
+    streamingSections[streamingSections.length - 1];
+  return { streamingSections, streamingSection, content };
+}
+
 export type AilaStreamingStatus =
   | "Loading"
   | "RequestMade"
@@ -18,36 +50,50 @@ 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";
+    log.info("Find streaming sections", lastMessage);
+    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/hooks/useProgressForDownloads.ts

@@ -1,6 +1,9 @@
 import { useMemo } from "react";
 
-import type { LooseLessonPlan } from "@oakai/aila/src/protocol/schema";
+import type {
+  LessonPlanKeys,
+  LooseLessonPlan,
+} from "@oakai/aila/src/protocol/schema";
 import { lessonPlanSectionsSchema } from "@oakai/exports/src/schema/input.schema";
 import type { ZodIssue } from "zod";
 
@@ -16,17 +19,20 @@ function getCompleteness(errors: ZodIssue[], fields: string[]) {
 
   return !hasErrorInSomeField;
 }
-export type ProgressSections = {
-  label: string;
-  key: string;
-  complete: boolean;
-}[];
+export type ProgressSections = ProgressSection[];
+
 type ProgressForDownloads = {
-  sections: ProgressSections;
+  sections: ProgressSection[];
   totalSections: number;
   totalSectionsComplete: number;
 };
 
+type ProgressSection = {
+  label: string;
+  key: LessonPlanKeys;
+  complete: boolean;
+};
+
 export function useProgressForDownloads({
   lessonPlan,
   isStreaming,
@@ -60,7 +66,8 @@ export function useProgressForDownloads({
           return true;
         }
       }) || [];
-    const sections = [
+
+    const sections: ProgressSection[] = [
       {
         label: "Lesson details",
         key: "title",