UI enhance: block rendering of historical messages until replay is co…

…mplete Signed-off-by: Ryan Hopper-Lowe <[email protected]>
obot-platform · Nov 5, 2024 · dbdbf2f · dbdbf2f
1 parent ea0e6a0
commit dbdbf2f
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 29 deletions.
diff --git a/ui/admin/app/components/chat/ChatContext.tsx b/ui/admin/app/components/chat/ChatContext.tsx
@@ -4,7 +4,6 @@ import {
     useCallback,
     useContext,
     useEffect,
-    useMemo,
     useState,
 } from "react";
 import { mutate } from "swr";
@@ -118,68 +117,78 @@ export function useChat() {
 }
 
 function useMessageSource(threadId?: string) {
-    const [messageMap, setMessageMap] = useState<Map<string, Message>>(
-        new Map()
-    );
+    const [messages, setMessages] = useState<Message[]>([]);
     const [isRunning, setIsRunning] = useState(false);
 
     const addContent = useCallback((event: ChatEvent) => {
-        console.log(event);
-
-        const { content, prompt, toolCall, runComplete, input, error, runID } =
-            event;
+        const {
+            content,
+            prompt,
+            toolCall,
+            runComplete,
+            input,
+            error,
+            runID,
+            contentID,
+        } = event;
 
         setIsRunning(!runComplete);
 
-        setMessageMap((prev) => {
-            const copy = new Map(prev);
+        setMessages((prev) => {
+            const copy = [...prev];
 
-            const contentID = event.contentID ?? crypto.randomUUID();
+            // todo(ryanhopperlowe) can be optmized by searching from the end
+            const existingIndex = contentID
+                ? copy.findIndex((m) => m.contentID === contentID)
+                : -1;
 
-            const existing = copy.get(contentID);
-            if (existing) {
-                copy.set(contentID, {
+            if (existingIndex !== -1) {
+                const existing = copy[existingIndex];
+                copy[existingIndex] = {
                     ...existing,
                     text: existing.text + content,
-                });
+                };
 
                 return copy;
             }
 
             if (error) {
-                copy.set(contentID, {
+                copy.push({
                     sender: "agent",
                     text: error,
                     runId: runID,
                     error: true,
+                    contentID,
                 });
                 return copy;
             }
 
             if (input) {
-                copy.set(contentID, {
+                copy.push({
                     sender: "user",
                     text: input,
                     runId: runID,
+                    contentID,
                 });
                 return copy;
             }
 
             if (toolCall) {
-                copy.set(contentID, toolCallMessage(toolCall));
+                copy.push(toolCallMessage(toolCall));
                 return copy;
             }
 
             if (prompt) {
-                copy.set(contentID, promptMessage(prompt, runID));
+                copy.push(promptMessage(prompt, runID));
                 return copy;
             }
 
             if (content) {
-                copy.set(contentID, {
+                copy.push({
                     sender: "agent",
                     text: content,
                     runId: runID,
+                    contentID,
                 });
                 return copy;
             }
@@ -189,25 +198,36 @@ function useMessageSource(threadId?: string) {
     }, []);
 
     useEffect(() => {
-        setMessageMap(new Map());
+        setMessages([]);
 
         if (!threadId) return;
 
         const source = ThreadsService.getThreadEventSource(threadId);
 
-        source.onmessage = (event) => {
-            const chunk = JSON.parse(event.data) as ChatEvent;
-            addContent(chunk);
+        let replayComplete = false;
+        let replayMessages: ChatEvent[] = [];
+
+        source.onmessage = (chunk) => {
+            const event = JSON.parse(chunk.data) as ChatEvent;
+
+            if (event.replayComplete) {
+                replayComplete = true;
+                replayMessages.forEach(addContent);
+                replayMessages = [];
+            }
+
+            if (!replayComplete) {
+                replayMessages.push(event);
+                return;
+            }
+
+            addContent(event);
         };
 
         return () => {
             source.close();
         };
     }, [threadId, addContent]);
 
-    const messages = useMemo(() => {
-        return Array.from(messageMap.values());
-    }, [messageMap]);
-
-    return { messages, messageMap, isRunning };
+    return { messages, isRunning };
 }
diff --git a/ui/admin/app/lib/model/chatEvents.ts b/ui/admin/app/lib/model/chatEvents.ts
@@ -38,6 +38,7 @@ export type ChatEvent = {
     content: string;
     input?: string;
     contentID?: string;
+    replayComplete?: boolean;
     error?: string;
     runComplete?: boolean;
     runID: string;

diff --git a/ui/admin/app/lib/model/messages.ts b/ui/admin/app/lib/model/messages.ts
@@ -11,6 +11,7 @@ export interface Message {
     runId?: string;
     isLoading?: boolean;
     error?: boolean;
+    contentID?: string;
 }
 
 export const runsToMessages = (runs: Run[]) => {