From ea0e6a0732999b5d834c220258752a2c1d387727 Mon Sep 17 00:00:00 2001 From: Ryan Hopper-Lowe Date: Tue, 5 Nov 2024 07:17:13 -0800 Subject: [PATCH] UI chore: load messages asynchronously via SSE Signed-off-by: Ryan Hopper-Lowe --- ui/admin/app/components/chat/Chat.tsx | 4 +- ui/admin/app/components/chat/ChatContext.tsx | 268 +++++++----------- ui/admin/app/components/chat/Chatbar.tsx | 7 +- ui/admin/app/components/chat/Message.tsx | 137 +++++---- ui/admin/app/components/chat/MessagePane.tsx | 20 +- ui/admin/app/components/chat/NoMessages.tsx | 5 +- ui/admin/app/lib/model/chatEvents.ts | 10 +- ui/admin/app/lib/routers/apiRoutes.ts | 7 +- .../app/lib/service/api/threadsService.ts | 22 ++ 9 files changed, 224 insertions(+), 256 deletions(-) diff --git a/ui/admin/app/components/chat/Chat.tsx b/ui/admin/app/components/chat/Chat.tsx index 175ce716f..fe88f94be 100644 --- a/ui/admin/app/components/chat/Chat.tsx +++ b/ui/admin/app/components/chat/Chat.tsx @@ -10,8 +10,7 @@ type ChatProps = React.HTMLAttributes & { }; export function Chat({ className, showStartButton = false }: ChatProps) { - const { messages, threadId, generatingMessage, mode, invoke, readOnly } = - useChat(); + const { messages, threadId, mode, invoke, readOnly } = useChat(); const [runTriggered, setRunTriggered] = useState(false); const showMessagePane = @@ -28,7 +27,6 @@ export function Chat({ className, showStartButton = false }: ChatProps) { )} diff --git a/ui/admin/app/components/chat/ChatContext.tsx b/ui/admin/app/components/chat/ChatContext.tsx index d909c788e..d1485cacb 100644 --- a/ui/admin/app/components/chat/ChatContext.tsx +++ b/ui/admin/app/components/chat/ChatContext.tsx @@ -1,25 +1,18 @@ import { ReactNode, createContext, - startTransition, + useCallback, useContext, useEffect, useMemo, - useRef, useState, } from "react"; -import useSWR, { mutate } from "swr"; +import { mutate } from "swr"; -import { ChatEvent, combineChatEvents } from "~/lib/model/chatEvents"; -import { - Message, - chatEventsToMessages, - promptMessage, - toolCallMessage, -} from "~/lib/model/messages"; +import { ChatEvent } from "~/lib/model/chatEvents"; +import { Message, promptMessage, toolCallMessage } from "~/lib/model/messages"; import { InvokeService } from "~/lib/service/api/invokeService"; import { ThreadsService } from "~/lib/service/api/threadsService"; -import { readStream } from "~/lib/stream"; import { useAsync } from "~/hooks/useAsync"; @@ -31,11 +24,10 @@ interface ChatContextType { processUserMessage: (text: string, sender: "user" | "agent") => void; id: string; threadId: string | undefined; - generatingMessage: Message | null; invoke: (prompt?: string) => void; readOnly?: boolean; isRunning: boolean; - isLoading: boolean; + isInvoking: boolean; } const ChatContext = createContext(undefined); @@ -55,64 +47,6 @@ export function ChatProvider({ onCreateThreadId?: (threadId: string) => void; readOnly?: boolean; }) { - const [insertedMessages, setInsertedMessages] = useState([]); - const [generatingMessage, setGeneratingMessage] = useState( - null - ); - const [isRunning, _setIsRunning] = useState(false); - const isRunningRef = useRef(false); - const isRunningToolCall = useRef(false); - - const setIsRunning = (value: boolean) => { - isRunningRef.current = value; - _setIsRunning(value); - }; - - // todo(tylerslaton): this is a huge hack to get the generating message and runId to be - // interactable during workflow invokes. take a look at invokeWorkflow to see why this is - // currently needed. - const generatingRunIdRef = useRef(null); - const generatingMessageRef = useRef(null); - - const appendToGeneratingMessage = (content: string) => { - generatingMessageRef.current = - (generatingMessageRef.current || "") + content; - - setGeneratingMessage(generatingMessageRef.current); - }; - - const clearGeneratingMessage = () => { - generatingMessageRef.current = null; - setGeneratingMessage(null); - }; - - const getThreadEvents = useSWR( - ThreadsService.getThreadEvents.key(threadId), - ({ threadId }) => ThreadsService.getThreadEvents(threadId), - { - onSuccess: () => setInsertedMessages([]), - revalidateOnFocus: false, - revalidateOnReconnect: false, - } - ); - - const messages = useMemo( - () => - chatEventsToMessages(combineChatEvents(getThreadEvents.data || [])), - [getThreadEvents.data] - ); - - // clear out inserted messages when the threadId changes - useEffect(() => { - if (isRunningRef.current) return; - setInsertedMessages([]); - }, [threadId]); - - /** inserts message optimistically */ - const insertMessage = (message: Message) => { - setInsertedMessages((prev) => [...prev, message]); - }; - /** * processUserMessage is responsible for adding the user's message to the chat and * triggering the agent to respond to it. @@ -121,7 +55,7 @@ export function ChatProvider({ if (mode === "workflow" || readOnly) return; const newMessage: Message = { text, sender }; - insertMessage(newMessage); + // insertMessage(newMessage); handlePrompt(newMessage.text); }; @@ -142,24 +76,8 @@ export function ChatProvider({ // do nothing if the mode is workflow }; - const insertGeneratingMessage = (runId?: string) => { - // skip if there is no message or it is only whitespace - if (generatingMessageRef.current) { - insertMessage({ - sender: "agent", - runId, - text: generatingMessageRef.current, - }); - clearGeneratingMessage(); - } - }; - const invokeAgent = useAsync(InvokeService.invokeAgentWithStream, { - onSuccess: ({ reader, threadId: responseThreadId }) => { - clearGeneratingMessage(); - - setIsRunning(true); - + onSuccess: ({ threadId: responseThreadId }) => { if (responseThreadId && !threadId) { // persist the threadId onCreateThreadId?.(responseThreadId); @@ -167,93 +85,22 @@ export function ChatProvider({ // revalidate threads mutate(ThreadsService.getThreads.key()); } - - readStream({ - reader, - onChunk: (chunk) => - // use a transition for performance - startTransition(() => { - const { content, toolCall, runID, input, prompt } = - chunk; - - generatingRunIdRef.current = runID; - - if (toolCall) { - isRunningToolCall.current = true; - // cut off generating message - insertGeneratingMessage(runID); - - // insert tool call message - insertMessage(toolCallMessage(toolCall)); - - clearGeneratingMessage(); - - return; - } - isRunningToolCall.current = false; - - if (prompt) { - insertGeneratingMessage(runID); - insertMessage(promptMessage(prompt, runID)); - return; - } - - if (content && !input) { - appendToGeneratingMessage(content); - } - }), - onComplete: async (chunks) => { - insertGeneratingMessage(chunks[0]?.runID); - clearGeneratingMessage(); - - invokeAgent.clear(); - generatingRunIdRef.current = null; - setIsRunning(false); - }, - }); }, }); - const outGeneratingMessage = useMemo(() => { - if (invokeAgent.isLoading) - return { sender: "agent", text: "", isLoading: true }; - - if (!generatingMessage) { - if (invokeAgent.data?.reader && !isRunningToolCall.current) { - return { - sender: "agent", - text: "", - isLoading: true, - }; - } - - return null; - } - - return { - sender: "agent", - text: generatingMessage, - runId: generatingRunIdRef.current ?? undefined, - }; - }, [generatingMessage, invokeAgent.isLoading, invokeAgent.data]); - - // combine messages and inserted messages - const outMessages = useMemo(() => { - return [...(messages ?? []), ...insertedMessages]; - }, [messages, insertedMessages]); + const { messages, isRunning } = useMessageSource(threadId); return ( @@ -269,3 +116,98 @@ export function useChat() { } return context; } + +function useMessageSource(threadId?: string) { + const [messageMap, setMessageMap] = useState>( + new Map() + ); + const [isRunning, setIsRunning] = useState(false); + + const addContent = useCallback((event: ChatEvent) => { + console.log(event); + + const { content, prompt, toolCall, runComplete, input, error, runID } = + event; + + setIsRunning(!runComplete); + + setMessageMap((prev) => { + const copy = new Map(prev); + + const contentID = event.contentID ?? crypto.randomUUID(); + + const existing = copy.get(contentID); + if (existing) { + copy.set(contentID, { + ...existing, + text: existing.text + content, + }); + + return copy; + } + + if (error) { + copy.set(contentID, { + sender: "agent", + text: error, + runId: runID, + error: true, + }); + return copy; + } + + if (input) { + copy.set(contentID, { + sender: "user", + text: input, + runId: runID, + }); + return copy; + } + + if (toolCall) { + copy.set(contentID, toolCallMessage(toolCall)); + return copy; + } + + if (prompt) { + copy.set(contentID, promptMessage(prompt, runID)); + return copy; + } + + if (content) { + copy.set(contentID, { + sender: "agent", + text: content, + runId: runID, + }); + return copy; + } + + return copy; + }); + }, []); + + useEffect(() => { + setMessageMap(new Map()); + + if (!threadId) return; + + const source = ThreadsService.getThreadEventSource(threadId); + + source.onmessage = (event) => { + const chunk = JSON.parse(event.data) as ChatEvent; + addContent(chunk); + }; + + return () => { + source.close(); + }; + }, [threadId, addContent]); + + const messages = useMemo(() => { + return Array.from(messageMap.values()); + }, [messageMap]); + + return { messages, messageMap, isRunning }; +} diff --git a/ui/admin/app/components/chat/Chatbar.tsx b/ui/admin/app/components/chat/Chatbar.tsx index c248eef48..cec3e9cea 100644 --- a/ui/admin/app/components/chat/Chatbar.tsx +++ b/ui/admin/app/components/chat/Chatbar.tsx @@ -4,6 +4,7 @@ import { useState } from "react"; import { cn } from "~/lib/utils"; import { useChat } from "~/components/chat/ChatContext"; +import { LoadingSpinner } from "~/components/ui/LoadingSpinner"; import { Button } from "~/components/ui/button"; import { AutosizeTextarea } from "~/components/ui/textarea"; @@ -13,7 +14,7 @@ type ChatbarProps = { export function Chatbar({ className }: ChatbarProps) { const [input, setInput] = useState(""); - const { processUserMessage, isRunning } = useChat(); + const { processUserMessage, isRunning, isInvoking } = useChat(); const handleSubmit = (e: React.FormEvent) => { e.preventDefault(); @@ -53,9 +54,9 @@ export function Chatbar({ className }: ChatbarProps) { variant="secondary" className="rounded-full" type="submit" - disabled={!input || isRunning} + disabled={!input || isRunning || isInvoking} > - + {isInvoking ? : } ); diff --git a/ui/admin/app/components/chat/Message.tsx b/ui/admin/app/components/chat/Message.tsx index afd5fa448..300acc2ee 100644 --- a/ui/admin/app/components/chat/Message.tsx +++ b/ui/admin/app/components/chat/Message.tsx @@ -16,10 +16,10 @@ import { CustomMarkdownComponents } from "~/components/react-markdown"; import { ToolIcon } from "~/components/tools/ToolIcon"; import { Button } from "~/components/ui/button"; import { Card } from "~/components/ui/card"; -import { TypingDots } from "~/components/ui/typing-spinner"; interface MessageProps { message: MessageType; + isRunning?: boolean; } // Allow links for file references in messages if it starts with file://, otherwise this will cause an empty href and cause app to reload when clicking on it @@ -54,80 +54,71 @@ export const Message = React.memo(({ message }: MessageProps) => {
- {message.isLoading ? ( - - ) : ( - +
+ {toolCall?.metadata?.icon && ( + + )} + + {message.prompt?.metadata ? ( + + ) : ( + + {parsedMessage || + "Waiting for more information..."} + + )} + + {toolCall && ( + + + )} - > -
- {toolCall?.metadata?.icon && ( - + - )} - - {message.prompt?.metadata ? ( - - ) : ( - - {parsedMessage || - "Waiting for more information..."} - - )} - - {toolCall && ( - - - - )} - - {!isUser && message.runId && ( -
- -
- )} - - {/* this is a hack to take up space for the debug button */} - {!toolCall && !message.runId && !isUser && ( -
-
- )} -
- - )} +
+ )} + + {/* this is a hack to take up space for the debug button */} + {!toolCall && !message.runId && !isUser && ( +
+
+ )} +
+ ); diff --git a/ui/admin/app/components/chat/MessagePane.tsx b/ui/admin/app/components/chat/MessagePane.tsx index 715d7f121..e03523840 100644 --- a/ui/admin/app/components/chat/MessagePane.tsx +++ b/ui/admin/app/components/chat/MessagePane.tsx @@ -5,8 +5,8 @@ import { cn } from "~/lib/utils"; import { useChat } from "~/components/chat/ChatContext"; import { Message } from "~/components/chat/Message"; import { NoMessages } from "~/components/chat/NoMessages"; -import { LoadingSpinner } from "~/components/ui/LoadingSpinner"; import { ScrollArea } from "~/components/ui/scroll-area"; +import { TypingDots } from "~/components/ui/typing-spinner"; interface MessagePaneProps { messages: MessageType[]; @@ -23,11 +23,10 @@ export function MessagePane({ messages, className, classNames = {}, - generatingMessage, }: MessagePaneProps) { - const { readOnly, isLoading } = useChat(); + const { readOnly, isRunning } = useChat(); - const isEmpty = messages.length === 0 && !generatingMessage && !readOnly; + const isEmpty = messages.length === 0 && !readOnly; return (
@@ -37,18 +36,19 @@ export function MessagePane({ enableScrollStick="bottom" className={cn("h-full w-full relative", classNames.messageList)} > - {isLoading && isEmpty ? ( - - ) : isEmpty ? ( + {isEmpty ? ( ) : (
{messages.map((message, i) => ( ))} - {generatingMessage && ( - - )} + +
)} diff --git a/ui/admin/app/components/chat/NoMessages.tsx b/ui/admin/app/components/chat/NoMessages.tsx index a39fb0366..d51285c24 100644 --- a/ui/admin/app/components/chat/NoMessages.tsx +++ b/ui/admin/app/components/chat/NoMessages.tsx @@ -4,7 +4,7 @@ import { useChat } from "~/components/chat/ChatContext"; import { Button } from "~/components/ui/button"; export function NoMessages() { - const { processUserMessage } = useChat(); + const { processUserMessage, isInvoking } = useChat(); const handleAddMessage = (content: string) => { processUserMessage(content, "user"); @@ -19,6 +19,7 @@ export function NoMessages() {