fix: update step duration display and logic (langflow-ai#4506)

* feat: Add durationStore for managing chat durations and intervals and update duration display logic to use store * [autofix.ci] apply automated fixes * Expand `recursive_serialize_or_str` to support `BaseModelV1` subclasses * fix: Update duration calculation for event handlers Improve accuracy of duration measurement in event handlers by centralizing the calculation method. This ensures consistent timing across different events and enhances the reliability of event processing. * refactor: improve duration calculation logic Enhance the duration calculation by clearly handling both integer and float timestamps. This ensures accurate duration tracking and simplifies the code structure, improving maintainability and readability. --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida <[email protected]> Co-authored-by: Cristhian Zanforlin Lousa <[email protected]>
headlinevc · Nov 26, 2024 · 78ca4fd · 78ca4fd
1 parent 471b28a
commit 78ca4fd
Show file tree

Hide file tree

Showing 7 changed files with 105 additions and 35 deletions.
diff --git a/src/backend/base/langflow/base/agents/events.py b/src/backend/base/langflow/base/agents/events.py
@@ -39,11 +39,17 @@ def _build_agent_input_text_content(agent_input_dict: InputDict) -> str:
 
 def _calculate_duration(start_time: float) -> int:
     """Calculate duration in milliseconds from start time to now."""
+    # Handle the calculation
+    current_time = perf_counter()
     if isinstance(start_time, int):
-        # means it was transformed into ms so we need to reverse it
-        # to whatever perf_counter returns
-        return int((perf_counter() - start_time / 1000) * 1000)
-    return int((perf_counter() - start_time) * 1000)
+        # If we got an integer, treat it as milliseconds
+        duration = current_time - (start_time / 1000)
+        result = int(duration * 1000)
+    else:
+        # If we got a float, treat it as perf_counter time
+        result = int((current_time - start_time) * 1000)
+
+    return result
 
 
 def handle_on_chain_start(
@@ -111,6 +117,9 @@ def handle_on_tool_start(
     if not agent_message.content_blocks:
         agent_message.content_blocks = [ContentBlock(title="Agent Steps", contents=[])]
 
+    duration = _calculate_duration(start_time)
+    new_start_time = perf_counter()  # Get new start time for next operation
+
     # Create new tool content with the input exactly as received
     tool_content = ToolContent(
         type="tool_use",
@@ -119,7 +128,7 @@ def handle_on_tool_start(
         output=None,
         error=None,
         header={"title": f"Accessing **{tool_name}**", "icon": "Hammer"},
-        duration=int(start_time * 1000),
+        duration=duration,  # Store the actual duration
     )
 
     # Store in map and append to message
@@ -128,7 +137,7 @@ def handle_on_tool_start(
 
     agent_message = send_message_method(message=agent_message)
     tool_blocks_map[tool_key] = agent_message.content_blocks[0].contents[-1]
-    return agent_message, start_time
+    return agent_message, new_start_time
 
 
 def handle_on_tool_end(
@@ -145,12 +154,13 @@ def handle_on_tool_end(
 
     if tool_content and isinstance(tool_content, ToolContent):
         tool_content.output = event["data"].get("output")
-        # Calculate duration only when tool ends
+        duration = _calculate_duration(start_time)
+        tool_content.duration = duration
         tool_content.header = {"title": f"Executed **{tool_content.name}**", "icon": "Hammer"}
-        if isinstance(tool_content.duration, int):
-            tool_content.duration = _calculate_duration(tool_content.duration)
+
         agent_message = send_message_method(message=agent_message)
-        start_time = perf_counter()
+        new_start_time = perf_counter()  # Get new start time for next operation
+        return agent_message, new_start_time
     return agent_message, start_time
 
 
@@ -250,12 +260,11 @@ async def process_agent_events(
                 agent_message, start_time = tool_handler(
                     event, agent_message, tool_blocks_map, send_message_method, start_time
                 )
-                start_time = start_time or perf_counter()
             elif event["event"] in CHAIN_EVENT_HANDLERS:
                 chain_handler = CHAIN_EVENT_HANDLERS[event["event"]]
                 agent_message, start_time = chain_handler(event, agent_message, send_message_method, start_time)
-                start_time = start_time or perf_counter()
         agent_message.properties.state = "complete"
     except Exception as e:
         raise ExceptionWithMessageError(agent_message) from e
+
     return Message(**agent_message.model_dump())
diff --git a/src/backend/base/langflow/schema/serialize.py b/src/backend/base/langflow/schema/serialize.py
@@ -8,7 +8,7 @@
 
 def recursive_serialize_or_str(obj):
     try:
-        if isinstance(obj, type) and issubclass(obj, BaseModel):
+        if isinstance(obj, type) and issubclass(obj, BaseModel | BaseModelV1):
             # This a type BaseModel and not an instance of it
             return repr(obj)
         if isinstance(obj, str):

diff --git a/src/frontend/src/components/chatComponents/ContentBlockDisplay.tsx b/src/frontend/src/components/chatComponents/ContentBlockDisplay.tsx
@@ -17,12 +17,14 @@ interface ContentBlockDisplayProps {
   contentBlocks: ContentBlock[];
   isLoading?: boolean;
   state?: string;
+  chatId: string;
 }
 
 export function ContentBlockDisplay({
   contentBlocks,
   isLoading,
   state,
+  chatId,
 }: ContentBlockDisplayProps) {
   const [isExpanded, setIsExpanded] = useState(false);
 
@@ -100,7 +102,7 @@ export function ContentBlockDisplay({
             </div>
           </div>
           <div className="flex items-center gap-2">
-            <DurationDisplay duration={totalDuration} />
+            <DurationDisplay duration={totalDuration} chatId={chatId} />
             <motion.div
               animate={{ rotate: isExpanded ? 180 : 0 }}
               transition={{ duration: 0.2, ease: "easeInOut" }}
@@ -190,7 +192,10 @@ export function ContentBlockDisplay({
                             </motion.div>
                           )}
                         </AnimatePresence>
-                        <ContentDisplay content={content} />
+                        <ContentDisplay
+                          content={content}
+                          chatId={`${chatId}-${index}`}
+                        />
                       </motion.div>
                     ))}
                   </div>

diff --git a/src/frontend/src/components/chatComponents/ContentDisplay.tsx b/src/frontend/src/components/chatComponents/ContentDisplay.tsx
@@ -8,7 +8,13 @@ import SimplifiedCodeTabComponent from "../codeTabsComponent/ChatCodeTabComponen
 import ForwardedIconComponent from "../genericIconComponent";
 import DurationDisplay from "./DurationDisplay";
 
-export default function ContentDisplay({ content }: { content: ContentType }) {
+export default function ContentDisplay({
+  content,
+  chatId,
+}: {
+  content: ContentType;
+  chatId: string;
+}) {
   // First render the common BaseContent elements if they exist
   const renderHeader = content.header && (
     <>
@@ -36,7 +42,7 @@ export default function ContentDisplay({ content }: { content: ContentType }) {
   );
   const renderDuration = content.duration !== undefined && (
     <div className="absolute right-2 top-4">
-      <DurationDisplay duration={content.duration} />
+      <DurationDisplay duration={content.duration} chatId={chatId} />
     </div>
   );
 

diff --git a/src/frontend/src/components/chatComponents/DurationDisplay.tsx b/src/frontend/src/components/chatComponents/DurationDisplay.tsx
@@ -1,34 +1,45 @@
-import { useEffect, useState } from "react";
+import { useDurationStore } from "@/stores/durationStore";
+import { useEffect } from "react";
 import { AnimatedNumber } from "../animatedNumbers";
 import ForwardedIconComponent from "../genericIconComponent";
 import Loading from "../ui/loading";
 
-export default function DurationDisplay({ duration }: { duration?: number }) {
-  const [elapsedTime, setElapsedTime] = useState(0);
-  const [intervalId, setIntervalId] = useState<NodeJS.Timeout | null>(null);
+interface DurationDisplayProps {
+  duration?: number;
+  chatId: string;
+}
+
+export default function DurationDisplay({
+  duration,
+  chatId,
+}: DurationDisplayProps) {
+  const {
+    durations,
+    setDuration,
+    incrementDuration,
+    clearInterval: clearDurationInterval,
+    setInterval: setDurationInterval,
+  } = useDurationStore();
 
   useEffect(() => {
-    if (duration !== undefined && intervalId) {
-      clearInterval(intervalId);
-      setIntervalId(null);
+    if (duration !== undefined) {
+      setDuration(chatId, duration);
+      clearDurationInterval(chatId);
       return;
     }
 
-    if (duration === undefined && !intervalId) {
-      const id = setInterval(() => {
-        setElapsedTime((prev) => prev + 10);
-      }, 10);
-      setIntervalId(id);
-    }
+    const intervalId = setInterval(() => {
+      incrementDuration(chatId);
+    }, 10);
+
+    setDurationInterval(chatId, intervalId);
 
     return () => {
-      if (intervalId) {
-        clearInterval(intervalId);
-      }
+      clearDurationInterval(chatId);
     };
-  }, [duration]);
+  }, [duration, chatId]);
 
-  const displayTime = duration ?? elapsedTime;
+  const displayTime = duration ?? durations[chatId] ?? 0;
   const secondsValue = displayTime / 1000;
   const humanizedTime = `${secondsValue.toFixed(1)}s`;
 

diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/newChatMessage.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/newChatMessage.tsx
@@ -479,6 +479,7 @@ export default function ChatMessage({
                   chat.properties?.state === "partial"
                 }
                 state={chat.properties?.state}
+                chatId={chat.id}
               />
             )}
             {!chat.isSend ? (

diff --git a/src/frontend/src/stores/durationStore.ts b/src/frontend/src/stores/durationStore.ts
@@ -0,0 +1,38 @@
+import { create } from "zustand";
+
+interface DurationState {
+  durations: Record<string, number>;
+  intervals: Record<string, NodeJS.Timeout>;
+  setDuration: (chatId: string, duration: number) => void;
+  incrementDuration: (chatId: string) => void;
+  clearInterval: (chatId: string) => void;
+  setInterval: (chatId: string, intervalId: NodeJS.Timeout) => void;
+}
+
+export const useDurationStore = create<DurationState>((set) => ({
+  durations: {},
+  intervals: {},
+  setDuration: (chatId, duration) =>
+    set((state) => ({
+      durations: { ...state.durations, [chatId]: duration },
+    })),
+  incrementDuration: (chatId) =>
+    set((state) => ({
+      durations: {
+        ...state.durations,
+        [chatId]: (state.durations[chatId] || 0) + 10,
+      },
+    })),
+  clearInterval: (chatId) =>
+    set((state) => {
+      if (state.intervals[chatId]) {
+        clearInterval(state.intervals[chatId]);
+      }
+      const { [chatId]: _, ...rest } = state.intervals;
+      return { intervals: rest };
+    }),
+  setInterval: (chatId, intervalId) =>
+    set((state) => ({
+      intervals: { ...state.intervals, [chatId]: intervalId },
+    })),
+}));