From 6d20913365fc2d520e632420014a580cc571d134 Mon Sep 17 00:00:00 2001
From: Ryan Hopper-Lowe <46546486+ryanhopperlowe@users.noreply.github.com>
Date: Fri, 18 Oct 2024 19:07:09 -0500
Subject: [PATCH] feat: split oauth prompts into their own messages (#222)

---
 ui/admin/app/components/chat/ChatContext.tsx | 11 ++++++--
 ui/admin/app/components/ui/scroll-area.tsx   |  4 +--
 ui/admin/app/lib/model/chatEvents.ts         | 21 ++++++++++++--
 ui/admin/app/lib/model/messages.ts           | 29 +++++++++++++-------
 4 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/ui/admin/app/components/chat/ChatContext.tsx b/ui/admin/app/components/chat/ChatContext.tsx
index 26c78d810..73cee7b84 100644
--- a/ui/admin/app/components/chat/ChatContext.tsx
+++ b/ui/admin/app/components/chat/ChatContext.tsx
@@ -14,6 +14,7 @@ import { ChatEvent, combineChatEvents } from "~/lib/model/chatEvents";
 import {
     Message,
     chatEventsToMessages,
+    promptMessage,
     toolCallMessage,
 } from "~/lib/model/messages";
 import { InvokeService } from "~/lib/service/api/invokeService";
@@ -149,7 +150,8 @@ export function ChatProvider({
                 onChunk: (chunk) =>
                     // use a transition for performance
                     startTransition(() => {
-                        const { content, toolCall, runID, input } = chunk;
+                        const { content, toolCall, runID, input, prompt } =
+                            chunk;
 
                         generatingRunIdRef.current = runID;
 
@@ -165,9 +167,14 @@ export function ChatProvider({
 
                             return;
                         }
-
                         isRunningToolCall.current = false;
 
+                        if (prompt) {
+                            insertGeneratingMessage(runID);
+                            insertMessage(promptMessage(prompt, runID));
+                            return;
+                        }
+
                         if (content && !input) {
                             appendToGeneratingMessage(content);
                         }
diff --git a/ui/admin/app/components/ui/scroll-area.tsx b/ui/admin/app/components/ui/scroll-area.tsx
index cb42f1d92..6f72a5669 100644
--- a/ui/admin/app/components/ui/scroll-area.tsx
+++ b/ui/admin/app/components/ui/scroll-area.tsx
@@ -40,8 +40,6 @@ const ScrollArea = React.forwardRef<
         }
     }, [startScrollAt]);
 
-    const deferredChildren = React.useDeferredValue(children);
-
     React.useEffect(() => {
         if (shouldStickToBottom && enableScrollStick === "bottom") {
             viewportRef.current?.scrollTo({
@@ -49,7 +47,7 @@ const ScrollArea = React.forwardRef<
                 behavior: "instant",
             });
         }
-    }, [shouldStickToBottom, enableScrollStick, deferredChildren]);
+    }, [enableScrollStick, shouldStickToBottom, children]);
 
     return (
         <ScrollAreaPrimitive.Root
diff --git a/ui/admin/app/lib/model/chatEvents.ts b/ui/admin/app/lib/model/chatEvents.ts
index 025a4c3cf..ae23ea0d0 100644
--- a/ui/admin/app/lib/model/chatEvents.ts
+++ b/ui/admin/app/lib/model/chatEvents.ts
@@ -13,6 +13,22 @@ export type ToolCall = {
     };
 };
 
+type PromptOAuthMeta = {
+    authType: string;
+    authURL: string;
+    toolContext: string;
+    toolDisplayName: string;
+};
+
+export type Prompt = {
+    id?: string;
+    time?: Date;
+    message?: string;
+    fields?: string[];
+    sensitive?: boolean;
+    metadata?: PromptOAuthMeta;
+};
+
 // note(ryanhopperlowe) renaming this to ChatEvent to differentiate itself specifically for a chat with an agent
 // we should create a separate type for WorkflowEvents and leverage Unions to differentiate between them
 export type ChatEvent = {
@@ -24,6 +40,7 @@ export type ChatEvent = {
     waitingOnModel?: boolean;
     toolInput?: ToolInput;
     toolCall?: ToolCall;
+    prompt?: Prompt;
 };
 
 export function combineChatEvents(events: ChatEvent[]): ChatEvent[] {
@@ -39,10 +56,10 @@ export function combineChatEvents(events: ChatEvent[]): ChatEvent[] {
     };
 
     for (const event of events) {
-        const { content, input, error, runID, toolCall } = event;
+        const { content, input, error, runID, toolCall, prompt } = event;
 
         // signals the end of a content block
-        if (error || toolCall || input) {
+        if (error || toolCall || input || prompt) {
             insertBuildingEvent();
 
             combinedEvents.push(event);
diff --git a/ui/admin/app/lib/model/messages.ts b/ui/admin/app/lib/model/messages.ts
index e6324be9a..cf6f61231 100644
--- a/ui/admin/app/lib/model/messages.ts
+++ b/ui/admin/app/lib/model/messages.ts
@@ -1,4 +1,4 @@
-import { ChatEvent, ToolCall } from "./chatEvents";
+import { ChatEvent, Prompt, ToolCall } from "./chatEvents";
 import { Run } from "./runs";
 
 export interface Message {
@@ -32,21 +32,24 @@ export const runsToMessages = (runs: Run[]) => {
     return messages;
 };
 
-export const toolCallMessage = (toolCall: ToolCall) => {
-    return {
-        sender: "agent",
-        text: `Tool call: ${[toolCall.metadata?.category, toolCall.name].filter((x) => !!x).join(" - ")}`,
-        tools: [toolCall],
-    } as Message;
-};
+export const toolCallMessage = (toolCall: ToolCall): Message => ({
+    sender: "agent",
+    text: `Tool call: ${[toolCall.metadata?.category, toolCall.name].filter((x) => !!x).join(" - ")}`,
+    tools: [toolCall],
+});
+
+export const promptMessage = (prompt: Prompt, runID: string): Message => ({
+    sender: "agent",
+    text: prompt.message || "",
+    runId: runID,
+});
 
 export const chatEventsToMessages = (events: ChatEvent[]) => {
     const messages: Message[] = [];
 
     for (const event of events) {
-        const { content, input, toolCall, runID, error } = event;
+        const { content, input, toolCall, runID, error, prompt } = event;
 
-        // skip errors and tool inputs with no content
         if (error) {
             messages.push({
                 sender: "agent",
@@ -71,6 +74,12 @@ export const chatEventsToMessages = (events: ChatEvent[]) => {
             continue;
         }
 
+        // note(ryanhopperlowe) this just splits out a new message. In the future we will want to create a custom prompt message
+        if (prompt) {
+            messages.push(promptMessage(prompt, runID));
+            continue;
+        }
+
         if (content) {
             messages.push({
                 sender: "agent",