From 9fd096c347f7dac6c6355a0e64a2084fe3307a78 Mon Sep 17 00:00:00 2001
From: Stanislas Polu <spolu@dust.tt>
Date: Wed, 27 Nov 2024 14:50:18 +0100
Subject: [PATCH] ActionsDetails: finer grained action state (#8940)

* ActionsDetails: finer grained action state

* nit

* memoize the spinner

* nit/comment

* tweak
---
 .../assistant/conversation/AgentMessage.tsx   | 17 ++++++
 .../actions/AgentMessageActions.tsx           | 61 ++++++++++++-------
 .../actions/AgentMessageActionsDrawer.tsx     |  6 +-
 3 files changed, 60 insertions(+), 24 deletions(-)

diff --git a/front/components/assistant/conversation/AgentMessage.tsx b/front/components/assistant/conversation/AgentMessage.tsx
index c28bb5db614d..b06f1b077d29 100644
--- a/front/components/assistant/conversation/AgentMessage.tsx
+++ b/front/components/assistant/conversation/AgentMessage.tsx
@@ -89,6 +89,8 @@ interface AgentMessageProps {
   size: ConversationMessageSizeType;
 }
 
+export type AgentStateClassification = "thinking" | "acting" | "done";
+
 /**
  *
  * @param isInModal is the conversation happening in a side modal, i.e. when
@@ -136,10 +138,18 @@ export function AgentMessage({
     }
   })();
 
+  const [lastAgentStateClassification, setLastAgentStateClassification] =
+    useState<AgentStateClassification>(shouldStream ? "thinking" : "done");
   const [lastTokenClassification, setLastTokenClassification] = useState<
     null | "tokens" | "chain_of_thought"
   >(null);
 
+  useEffect(() => {
+    if (message.status !== "created") {
+      setLastAgentStateClassification("done");
+    }
+  }, [message.status]);
+
   const buildEventSourceURL = useCallback(
     (lastEvent: string | null) => {
       const esURL = `/api/w/${owner.sId}/assistant/conversations/${conversationId}/messages/${message.sId}/events`;
@@ -187,6 +197,7 @@ export function AgentMessage({
         setStreamedAgentMessage((m) => {
           return { ...updateMessageWithAction(m, event.action) };
         });
+        setLastAgentStateClassification("thinking");
         break;
       case "retrieval_params":
       case "dust_app_run_params":
@@ -201,17 +212,20 @@ export function AgentMessage({
         setStreamedAgentMessage((m) => {
           return updateMessageWithAction(m, event.action);
         });
+        setLastAgentStateClassification("acting");
         break;
       case "agent_error":
         setStreamedAgentMessage((m) => {
           return { ...m, status: "failed", error: event.error };
         });
+        setLastAgentStateClassification("done");
         break;
 
       case "agent_generation_cancelled":
         setStreamedAgentMessage((m) => {
           return { ...m, status: "cancelled" };
         });
+        setLastAgentStateClassification("done");
         break;
       case "agent_message_success": {
         setStreamedAgentMessage((m) => {
@@ -220,6 +234,7 @@ export function AgentMessage({
             ...event.message,
           };
         });
+        setLastAgentStateClassification("done");
         break;
       }
 
@@ -249,6 +264,7 @@ export function AgentMessage({
           default:
             assertNever(event);
         }
+        setLastAgentStateClassification("thinking");
         break;
       }
 
@@ -528,6 +544,7 @@ export function AgentMessage({
       <div className="flex flex-col gap-y-4">
         <AgentMessageActions
           agentMessage={agentMessage}
+          lastAgentStateClassification={lastAgentStateClassification}
           size={size}
           owner={owner}
         />
diff --git a/front/components/assistant/conversation/actions/AgentMessageActions.tsx b/front/components/assistant/conversation/actions/AgentMessageActions.tsx
index 014eee8653f7..97ef46604c51 100644
--- a/front/components/assistant/conversation/actions/AgentMessageActions.tsx
+++ b/front/components/assistant/conversation/actions/AgentMessageActions.tsx
@@ -5,54 +5,60 @@ import type {
   AgentMessageType,
   LightWorkspaceType,
 } from "@dust-tt/types";
+import { assertNever } from "@dust-tt/types";
 import { useEffect, useMemo, useState } from "react";
 
 import { getActionSpecification } from "@app/components/actions/types";
 import { AgentMessageActionsDrawer } from "@app/components/assistant/conversation/actions/AgentMessageActionsDrawer";
+import type { AgentStateClassification } from "@app/components/assistant/conversation/AgentMessage";
 import { classNames } from "@app/lib/utils";
+
 interface AgentMessageActionsProps {
   agentMessage: AgentMessageType;
+  lastAgentStateClassification: AgentStateClassification;
   size?: ConversationMessageSizeType;
   owner: LightWorkspaceType;
 }
 
 export function AgentMessageActions({
   agentMessage,
+  lastAgentStateClassification,
   owner,
   size = "normal",
 }: AgentMessageActionsProps) {
   const [chipLabel, setChipLabel] = useState<string | undefined>("Thinking");
   const [isActionDrawerOpened, setIsActionDrawerOpened] = useState(false);
 
-  // We're thinking or acting if the message status is still "created" and we don't have content
-  // yet. Despite our work on chain of thoughts events, it's still possible for content to be
-  // emitted before actions in which case we will think we're not thinking or acting until an action
-  // gets emitted in which case the content will get requalified as chain of thoughts and this will
-  // switch back to true.
+  useEffect(() => {
+    switch (lastAgentStateClassification) {
+      case "thinking":
+        setChipLabel("Thinking");
+        break;
+      case "acting":
+        if (agentMessage.actions.length > 0) {
+          setChipLabel(renderActionName(agentMessage.actions));
+        }
+        break;
+      case "done":
+        setChipLabel(undefined);
+        break;
+      default:
+        assertNever(lastAgentStateClassification);
+    }
+  }, [lastAgentStateClassification, agentMessage.actions]);
+
   const isThinkingOrActing = useMemo(
     () => agentMessage.status === "created",
     [agentMessage.status]
   );
 
-  useEffect(() => {
-    if (isThinkingOrActing) {
-      if (agentMessage.actions.length === 0) {
-        setChipLabel("Thinking");
-      } else {
-        setChipLabel(renderActionName(agentMessage.actions));
-      }
-    } else {
-      setChipLabel(undefined);
-    }
-  }, [isThinkingOrActing, agentMessage.actions]);
-
   return (
     <div className="flex flex-col items-start gap-y-4">
       <AgentMessageActionsDrawer
         actions={agentMessage.actions}
         isOpened={isActionDrawerOpened}
         onClose={() => setIsActionDrawerOpened(false)}
-        isStreaming={isThinkingOrActing || agentMessage.actions.length === 0}
+        isActing={lastAgentStateClassification === "acting"}
         owner={owner}
       />
       <ActionDetails
@@ -79,12 +85,19 @@ function ActionDetails({
   onClick: () => void;
   size: ConversationMessageSizeType;
 }) {
+  // We memoize the spinner as otherwise its state gets resetted on each token emission (despite
+  // memoization of label in the parent component).
+  const MemoizedSpinner = useMemo(
+    () => <Spinner variant="dark" size="xs" />,
+    []
+  );
+
   if (!label && (!isActionStepDone || !hasActions)) {
     return null;
   }
 
   return label ? (
-    <div key={label} className="animate-fadeIn duration-1000 fade-out">
+    <div key={label}>
       <Chip size="sm" color="slate" isBusy>
         <div
           className={classNames(
@@ -93,8 +106,14 @@ function ActionDetails({
           )}
           onClick={hasActions ? onClick : undefined}
         >
-          <Spinner variant="dark" size="xs" />
-          {label}
+          {MemoizedSpinner}
+          {label === "Thinking" ? (
+            <span>{label}</span>
+          ) : (
+            <>
+              Thinking <span className="text-regular">{label}</span>
+            </>
+          )}
         </div>
       </Chip>
     </div>
diff --git a/front/components/assistant/conversation/actions/AgentMessageActionsDrawer.tsx b/front/components/assistant/conversation/actions/AgentMessageActionsDrawer.tsx
index 430a11b9ad5b..95009f878bdf 100644
--- a/front/components/assistant/conversation/actions/AgentMessageActionsDrawer.tsx
+++ b/front/components/assistant/conversation/actions/AgentMessageActionsDrawer.tsx
@@ -6,7 +6,7 @@ import { getActionSpecification } from "@app/components/actions/types";
 interface AgentMessageActionsDrawerProps {
   actions: AgentActionType[];
   isOpened: boolean;
-  isStreaming: boolean;
+  isActing: boolean;
   onClose: () => void;
   owner: LightWorkspaceType;
 }
@@ -14,7 +14,7 @@ interface AgentMessageActionsDrawerProps {
 export function AgentMessageActionsDrawer({
   actions,
   isOpened,
-  isStreaming,
+  isActing,
   onClose,
   owner,
 }: AgentMessageActionsDrawerProps) {
@@ -71,7 +71,7 @@ export function AgentMessageActionsDrawer({
                 </div>
               );
             })}
-            {isStreaming && (
+            {isActing && (
               <div className="flex justify-center">
                 <Spinner variant="color" />
               </div>