From 1b6339290399276ee3b4f78cb7d078765d219195 Mon Sep 17 00:00:00 2001 From: Ryan Hopper-Lowe Date: Wed, 15 Jan 2025 08:42:43 -0600 Subject: [PATCH] enhance: improve tool authentication dialog UX --- .../agent/shared/ToolAuthenticationDialog.tsx | 161 ++++++++++++--- .../agent/shared/ToolAuthenticationStatus.tsx | 1 - ui/admin/app/components/chat/ChatContext.tsx | 189 +----------------- ui/admin/app/components/chat/Message.tsx | 26 ++- .../app/hooks/messages/useMessageSource.ts | 161 +++++++++++++++ 5 files changed, 316 insertions(+), 222 deletions(-) create mode 100644 ui/admin/app/hooks/messages/useMessageSource.ts diff --git a/ui/admin/app/components/agent/shared/ToolAuthenticationDialog.tsx b/ui/admin/app/components/agent/shared/ToolAuthenticationDialog.tsx index 164024c32..23afe22dd 100644 --- a/ui/admin/app/components/agent/shared/ToolAuthenticationDialog.tsx +++ b/ui/admin/app/components/agent/shared/ToolAuthenticationDialog.tsx @@ -1,9 +1,11 @@ -import { useCallback, useState } from "react"; +import { CheckIcon, CircleAlert } from "lucide-react"; +import { useEffect, useMemo, useState } from "react"; -import { ChatEvent } from "~/lib/model/chatEvents"; +import { ThreadsService } from "~/lib/service/api/threadsService"; import { useToolReference } from "~/components/agent/ToolEntry"; -import { Chat, ChatProvider } from "~/components/chat"; +import { PromptAuthForm } from "~/components/chat/Message"; +import { LoadingSpinner } from "~/components/ui/LoadingSpinner"; import { Button } from "~/components/ui/button"; import { Dialog, @@ -14,35 +16,59 @@ import { DialogHeader, DialogTitle, } from "~/components/ui/dialog"; +import { Link } from "~/components/ui/link"; +import { useMessageStream } from "~/hooks/messages/useMessageSource"; type AgentAuthenticationDialogProps = { threadId: Nullish; onComplete: () => void; - entityId: string; tool: string; }; export function ToolAuthenticationDialog({ onComplete, threadId, - entityId, tool, }: AgentAuthenticationDialogProps) { - const [done, setDone] = useState(false); - const handleDone = useCallback(() => setDone(true), []); - const { icon, label } = useToolReference(tool); - const onRunEvent = useCallback( - ({ content }: ChatEvent) => { - if (content === "DONE") handleDone(); - }, - [handleDone] + const source = useMemo( + () => (threadId ? ThreadsService.getThreadEventSource(threadId) : null), + [threadId] ); + const { messages: _messages } = useMessageStream(source); + + type ItemState = { + isLoading?: boolean; + isError?: boolean; + isDone?: boolean; + }; + + const [map, setMap] = useState>({}); + const updateItem = (id: number, state: Partial) => + setMap((prev) => ({ ...prev, [id]: { ...prev[id], ...state } })); + + const messages = useMemo( + () => _messages.filter((m) => m.prompt || m.error || m.text === "DONE"), + [_messages] + ); + + useEffect(() => { + // any time a message is added, prevent the last message from being loading + const isError = messages.at(-1)?.error; + + const i = messages.length - 2; + setMap((prev) => ({ + ...prev, + [i]: { isLoading: false, isDone: !isError, isError }, + })); + }, [messages]); + + const done = messages.at(-1)?.text === "DONE"; return ( - + {icon} Authorize {label} @@ -51,25 +77,94 @@ export function ToolAuthenticationDialog({ - {done ? ( - - {label} has successfully been authorized. You may now close this - modal. - - ) : ( - - - - )} +
+ {!messages.length ? ( +
+ Loading... +
+ ) : ( +
+ {messages.map((message, index) => { + if (message.error) { + return ( +

+ Error: {message.text} +

+ ); + } + + if (message.text === "DONE") { + return ( +

+ + Done +

+ ); + } + + if (message.prompt) { + if (map[index]?.isDone) { + return ( +

+ + Authentication Successful +

+ ); + } + + if (map[index]?.isLoading) { + return ( +

+ Authentication Processing +

+ ); + } + + if (message.prompt.metadata?.authURL) { + return ( +

+ + + Authentication Required{" "} + + updateItem(index, { isLoading: true }) + } + > + Click Here + + +

+ ); + } + + if (message.prompt.fields) { + return ( +
+

+ + Authentication Required +

+ + updateItem(index, { isLoading: true }) + } + /> +
+ ); + } + } + })} +
+ )} +
diff --git a/ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx b/ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx index 90acb8df4..eca719fd5 100644 --- a/ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx +++ b/ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx @@ -160,7 +160,6 @@ export function ToolAuthenticationStatus({ diff --git a/ui/admin/app/components/chat/ChatContext.tsx b/ui/admin/app/components/chat/ChatContext.tsx index 1c2ddf89f..d495e9b3b 100644 --- a/ui/admin/app/components/chat/ChatContext.tsx +++ b/ui/admin/app/components/chat/ChatContext.tsx @@ -1,18 +1,11 @@ -import { - ReactNode, - createContext, - useCallback, - useContext, - useEffect, - useState, -} from "react"; +import { ReactNode, createContext, useContext, useMemo } from "react"; import { mutate } from "swr"; -import { ChatEvent } from "~/lib/model/chatEvents"; -import { Message, promptMessage, toolCallMessage } from "~/lib/model/messages"; +import { Message } from "~/lib/model/messages"; import { InvokeService } from "~/lib/service/api/invokeService"; import { ThreadsService } from "~/lib/service/api/threadsService"; +import { useMessageStream } from "~/hooks/messages/useMessageSource"; import { useAsync } from "~/hooks/useAsync"; type Mode = "agent" | "workflow"; @@ -39,7 +32,6 @@ export function ChatProvider({ threadId, onCreateThreadId, readOnly, - onRunEvent, }: { children: ReactNode; mode?: Mode; @@ -47,8 +39,6 @@ export function ChatProvider({ threadId?: Nullish; onCreateThreadId?: (threadId: string) => void; readOnly?: boolean; - /** @description THIS MUST BE MEMOIZED */ - onRunEvent?: (event: ChatEvent) => void; }) { const invoke = (prompt?: string) => { if (readOnly) return; @@ -70,7 +60,12 @@ export function ChatProvider({ }, }); - const { messages, isRunning } = useMessageSource(threadId, onRunEvent); + const source = useMemo( + () => (threadId ? ThreadsService.getThreadEventSource(threadId) : null), + [threadId] + ); + + const { messages, isRunning } = useMessageStream(source); const abortRunningThread = () => { if (!threadId || !isRunning) return; @@ -106,169 +101,3 @@ export function useChat() { } return context; } - -function useMessageSource( - threadId?: Nullish, - onRunEvent?: (event: ChatEvent) => void -) { - const [messages, setMessages] = useState([]); - const [isRunning, setIsRunning] = useState(false); - - const addContent = useCallback( - (event: ChatEvent) => { - const { - content, - prompt, - toolCall, - runComplete, - input, - error, - runID, - contentID, - replayComplete, - } = event; - - onRunEvent?.(event); - - setIsRunning(!runComplete && !replayComplete); - - setMessages((prev) => { - const copy = [...prev]; - - // todo(ryanhopperlowe) can be optmized by searching from the end - const existingIndex = contentID - ? copy.findIndex((m) => m.contentID === contentID) - : -1; - - if (existingIndex !== -1) { - const existing = copy[existingIndex]; - copy[existingIndex] = { - ...existing, - text: existing.text + content, - }; - - return copy; - } - - if (error) { - if (error.includes("thread was aborted, cancelling run")) { - copy.push({ - sender: "agent", - text: "Message Aborted", - runId: runID, - contentID, - aborted: true, - }); - - return copy; - } - - copy.push({ - sender: "agent", - text: error, - runId: runID, - error: true, - contentID, - }); - return copy; - } - - if (input) { - copy.push({ - sender: "user", - text: input, - runId: runID, - contentID, - }); - return copy; - } - - if (toolCall) { - return handleToolCallEvent(copy, event); - } - - if (prompt) { - copy.push(promptMessage(prompt, runID)); - return copy; - } - - if (content) { - copy.push({ - sender: "agent", - text: content, - runId: runID, - contentID, - }); - return copy; - } - - return copy; - }); - }, - [onRunEvent] - ); - - useEffect(() => { - setMessages([]); - - if (!threadId) return; - - let replayComplete = false; - let replayMessages: ChatEvent[] = []; - - const source = ThreadsService.getThreadEventSource(threadId); - source.addEventListener("close", source.close); - - 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(); - setIsRunning(false); - }; - }, [threadId, addContent]); - - return { messages, isRunning }; -} - -const findIndexLastPendingToolCall = (messages: Message[]) => { - for (let i = messages.length - 1; i >= 0; i--) { - const message = messages[i]; - if (message.tools && !message.tools[0].output) { - return i; - } - } - return null; -}; - -const handleToolCallEvent = (messages: Message[], event: ChatEvent) => { - if (!event.toolCall) return messages; - - const { toolCall } = event; - if (toolCall.output) { - const index = findIndexLastPendingToolCall(messages); - if (index !== null) { - // update the found pending toolcall message (without output) - messages[index].tools = [toolCall]; - return messages; - } - } - - // otherwise add a new toolcall message - messages.push(toolCallMessage(toolCall)); - return messages; -}; diff --git a/ui/admin/app/components/chat/Message.tsx b/ui/admin/app/components/chat/Message.tsx index 2761ec95e..ba6e5abbc 100644 --- a/ui/admin/app/components/chat/Message.tsx +++ b/ui/admin/app/components/chat/Message.tsx @@ -52,6 +52,8 @@ export const Message = React.memo(({ message }: MessageProps) => { // leaving it in case that changes in the future const [toolCall = null] = message.tools || []; + const { isRunning } = useChat(); + const parsedMessage = useMemo(() => { if (OpenMarkdownLinkRegex.test(message.text)) { return message.text.replace( @@ -89,7 +91,7 @@ export const Message = React.memo(({ message }: MessageProps) => { )} {message.prompt ? ( - + ) : ( { Message.displayName = "Message"; -function PromptMessage({ prompt }: { prompt: AuthPrompt }) { +export function PromptMessage({ + prompt, + isRunning = false, +}: { + prompt: AuthPrompt; + isRunning?: boolean; +}) { const [open, setOpen] = useState(false); const [isSubmitted, setIsSubmitted] = useState(false); - const { isRunning } = useChat(); const getMessage = () => { if (prompt.metadata?.authURL || prompt.metadata?.authType) @@ -272,12 +279,14 @@ function PromptMessage({ prompt }: { prompt: AuthPrompt }) { ); } -function PromptAuthForm({ +export function PromptAuthForm({ prompt, onSuccess, + onSubmit, }: { prompt: AuthPrompt; - onSuccess: () => void; + onSuccess?: () => void; + onSubmit?: () => void; }) { const authenticate = useAsync(PromptApiService.promptResponse, { onSuccess, @@ -293,9 +302,10 @@ function PromptAuthForm({ ), }); - const handleSubmit = form.handleSubmit(async (values) => - authenticate.execute({ id: prompt.id, response: values }) - ); + const handleSubmit = form.handleSubmit(async (values) => { + authenticate.execute({ id: prompt.id, response: values }); + onSubmit?.(); + }); return (
diff --git a/ui/admin/app/hooks/messages/useMessageSource.ts b/ui/admin/app/hooks/messages/useMessageSource.ts new file mode 100644 index 000000000..970c83d68 --- /dev/null +++ b/ui/admin/app/hooks/messages/useMessageSource.ts @@ -0,0 +1,161 @@ +import { useCallback, useEffect, useState } from "react"; + +import { ChatEvent } from "~/lib/model/chatEvents"; +import { Message, promptMessage, toolCallMessage } from "~/lib/model/messages"; + +export function useMessageStream(source: Nullish) { + const [messages, setMessages] = useState([]); + const [isRunning, setIsRunning] = useState(false); + + const addContent = useCallback((event: ChatEvent) => { + const { + content, + prompt, + toolCall, + runComplete, + input, + error, + runID, + contentID, + replayComplete, + } = event; + + setIsRunning(!runComplete && !replayComplete); + + setMessages((prev) => { + const copy = [...prev]; + + // todo(ryanhopperlowe) can be optmized by searching from the end + const existingIndex = contentID + ? copy.findIndex((m) => m.contentID === contentID) + : -1; + + if (existingIndex !== -1) { + const existing = copy[existingIndex]; + copy[existingIndex] = { + ...existing, + text: existing.text + content, + }; + + return copy; + } + + if (error) { + if (error.includes("thread was aborted, cancelling run")) { + copy.push({ + sender: "agent", + text: "Message Aborted", + runId: runID, + contentID, + aborted: true, + }); + + return copy; + } + + copy.push({ + sender: "agent", + text: error, + runId: runID, + error: true, + contentID, + }); + return copy; + } + + if (input) { + copy.push({ + sender: "user", + text: input, + runId: runID, + contentID, + }); + return copy; + } + + if (toolCall) { + return handleToolCallEvent(copy, event); + } + + if (prompt) { + copy.push(promptMessage(prompt, runID)); + return copy; + } + + if (content) { + copy.push({ + sender: "agent", + text: content, + runId: runID, + contentID, + }); + return copy; + } + + return copy; + }); + }, []); + + useEffect(() => { + setMessages([]); + + if (!source) return; + + let replayComplete = false; + let replayMessages: ChatEvent[] = []; + + source.addEventListener("close", source.close); + + source.addEventListener("message", (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(); + setIsRunning(false); + }; + }, [source, addContent]); + + return { messages, isRunning }; +} + +const findIndexLastPendingToolCall = (messages: Message[]) => { + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i]; + if (message.tools && !message.tools[0].output) { + return i; + } + } + return null; +}; + +const handleToolCallEvent = (messages: Message[], event: ChatEvent) => { + if (!event.toolCall) return messages; + + const { toolCall } = event; + if (toolCall.output) { + const index = findIndexLastPendingToolCall(messages); + if (index !== null) { + // update the found pending toolcall message (without output) + messages[index].tools = [toolCall]; + return messages; + } + } + + // otherwise add a new toolcall message + messages.push(toolCallMessage(toolCall)); + return messages; +};