display errors from the backend (#425)

Signed-off-by: Peter Jausovec <peter.jausovec@solo.io>
Co-authored-by: Eitan Yarmush <eitan.yarmush@solo.io>

Author: peterj

python/src/autogenstudio/sessionmanager/sessionmanager.py

@@ -167,10 +167,14 @@ async def start_stream(
                 logger.error(f"Stream error for run {run_id}: {e}")
                 traceback.print_exc()
 
+                # The messages[0].content isn't properly being serialized, so it
+                # doesn't even get sent back to the client. (That's why we're seeing undefined in the UI)
+                # I am using the stop_reason to send the error message back and specifically checking 
+                # for the message type (error).
                 error_result = TeamResult(
                     task_result=TaskResult(
                         messages=[TextMessage(source="system", content=str(e))],
-                        stop_reason="An error occurred while processing this run",
+                        stop_reason=str(e),
                     ),
                     usage="",
                     duration=0,

ui/src/components/chat/ChatInterface.tsx

@@ -5,18 +5,18 @@ import { useState, useRef, useEffect } from "react";
 import { ArrowBigUp, X, Loader2 } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Textarea } from "@/components/ui/textarea";
-import type { Session, AgentMessageConfig, TextMessageConfig } from "@/types/datamodel";
+import type { Session, AgentMessageConfig } from "@/types/datamodel";
 import { ScrollArea } from "@/components/ui/scroll-area";
 import ChatMessage from "@/components/chat/ChatMessage";
 import StreamingMessage from "./StreamingMessage";
-import TokenStatsDisplay, { calculateTokenStats } from "./TokenStats";
+import TokenStatsDisplay from "./TokenStats";
 import { TokenStats } from "@/lib/types";
 import StatusDisplay from "./StatusDisplay";
 import { createSession, getSessionMessages, checkSessionExists, updateSession } from "@/app/actions/sessions";
 import { getCurrentUserId } from "@/app/actions/utils";
-import { messageUtils } from "@/lib/utils";
 import { toast } from "sonner";
 import { useRouter } from "next/navigation";
+import { createMessageHandlers } from "@/lib/messageHandlers";
 
 export type ChatStatus = "ready" | "thinking" | "error";
 
@@ -49,6 +49,13 @@ export default function ChatInterface({ selectedAgentId, selectedSession, sessio
   const isCreatingSessionRef = useRef<boolean>(false);
   const [isFirstMessage, setIsFirstMessage] = useState<boolean>(!sessionId);
 
+  const { handleMessageEvent } = createMessageHandlers({
+    setMessages,
+    setIsStreaming,
+    setStreamingContent,
+    setTokenStats
+  });
+
   useEffect(() => {
     async function initializeChat() {
       // Skip completely if this is a first message session creation flow
@@ -251,37 +258,7 @@ export default function ChatInterface({ selectedAgentId, selectedSession, sessio
               if (eventData) {
                 try {
                   const eventDataJson = JSON.parse(eventData) as AgentMessageConfig;
-
-                  if (messageUtils.isStreamingMessage(eventDataJson)) {
-                    // Set the streaming flag to true and concatenate the content
-                    setIsStreaming(true);
-                    setStreamingContent(prev => prev + eventDataJson.content);
-                  } else if (messageUtils.isTextMessageContent(eventDataJson)) {
-                    // The model usage is sent within the TextMessage, after the streaming is complete
-                    setTokenStats(prev => calculateTokenStats(prev, eventDataJson as TextMessageConfig));
-                    setIsStreaming(false);
-                    setStreamingContent("");
-                    if (eventDataJson.source !== "user") {
-                      // We don't want to add the user's message to the messages array (again), because 
-                      // we already added it when the user sent the message.
-                      setMessages(prevMessages => [...prevMessages, eventDataJson]);
-                    }
-                  }
-                  else {
-                    // For tool call messages and other non-streaming messages,
-                    // add them to the messages array immediately but don't stop streaming
-                    if (messageUtils.isToolCallRequestEvent(eventDataJson) || 
-                        messageUtils.isToolCallExecutionEvent(eventDataJson) || 
-                        messageUtils.isToolCallSummaryMessage(eventDataJson)) {
-                      // Add tool call messages immediately without stopping streaming
-                      setMessages(prevMessages => [...prevMessages, eventDataJson]);
-                    } else {
-                      // For other non-tool, non-streaming messages, use original behavior
-                      setIsStreaming(false);
-                      setStreamingContent("");
-                      setMessages(prevMessages => [...prevMessages, eventDataJson]);
-                    }
-                  }
+                  handleMessageEvent(eventDataJson);
                 } catch (error) {
                   toast.error("Error parsing event data");
                   console.error("Error parsing event data:", error, eventData);

ui/src/components/chat/ChatMessage.tsx

@@ -25,7 +25,12 @@ export default function ChatMessage({ message, allMessages }: ChatMessageProps)
     return <ToolCallDisplay currentMessage={message} allMessages={allMessages} />;
   }
 
-  const { content, source } = message;
+  let { content, source } = message;
+
+  const isErrorMessage = messageUtils.isErrorMessageContent(message);
+  if (isErrorMessage) {
+    content = message.data.task_result.stop_reason || "An error occurred";
+  }
 
   // Filter out system messages
   // TODO: Decide whether we want to filter out som agent
@@ -42,8 +47,9 @@ export default function ChatMessage({ message, allMessages }: ChatMessageProps)
     return <LLMCallModal content={String(message)} />;
   }
 
+  const messageBorderColor = isErrorMessage ? "border-l-red-500" : source === "user" ? "border-l-blue-500" : "border-l-violet-500";
 
-  return <div className={`flex items-center gap-2 text-sm border-l-2 py-2 px-4 ${source === "user" ? "border-l-blue-500" : "border-l-violet-500"}`}>
+  return <div className={`flex items-center gap-2 text-sm border-l-2 py-2 px-4 ${messageBorderColor}`}>
     <div className="flex flex-col gap-1 w-full">
       {source !== "user" ? <div className="flex items-center gap-1">
         <KagentLogo className="w-4 h-4" />

ui/src/lib/messageHandlers.ts

@@ -0,0 +1,75 @@
+import { AgentMessageConfig, TextMessageConfig } from "@/types/datamodel";
+import { messageUtils } from "@/lib/utils";
+import { TokenStats } from "@/lib/types";
+import { calculateTokenStats } from "@/components/chat/TokenStats";
+
+export type MessageHandlers = {
+  setMessages: (updater: (prev: AgentMessageConfig[]) => AgentMessageConfig[]) => void;
+  setIsStreaming: (value: boolean) => void;
+  setStreamingContent: (updater: (prev: string) => string) => void;
+  setTokenStats: (updater: (prev: TokenStats) => TokenStats) => void;
+};
+
+export const createMessageHandlers = (handlers: MessageHandlers) => {
+  const handleStreamingMessage = (message: AgentMessageConfig) => {
+    handlers.setIsStreaming(true);
+    handlers.setStreamingContent(prev => prev + message.content);
+  };
+
+  const handleTextMessage = (message: AgentMessageConfig) => {
+    handlers.setTokenStats(prev => calculateTokenStats(prev, message as TextMessageConfig));
+    handlers.setIsStreaming(false);
+    handlers.setStreamingContent(() => "");
+    
+    // Only add non-user messages to the messages array
+    if (message.source !== "user") {
+      handlers.setMessages(prevMessages => [...prevMessages, message]);
+    }
+  };
+
+  const handleErrorMessage = (message: AgentMessageConfig) => {
+    handlers.setMessages(prevMessages => [...prevMessages, message]);
+  };
+
+  const handleToolCallMessage = (message: AgentMessageConfig) => {
+    handlers.setMessages(prevMessages => [...prevMessages, message]);
+  };
+
+  const handleOtherMessage = (message: AgentMessageConfig) => {
+    handlers.setIsStreaming(false);
+    handlers.setStreamingContent(() => "");
+    handlers.setMessages(prevMessages => [...prevMessages, message]);
+  };
+
+  const handleMessageEvent = (message: AgentMessageConfig) => {
+    if (messageUtils.isStreamingMessage(message)) {
+      handleStreamingMessage(message);
+      return;
+    }
+
+    if (messageUtils.isTextMessageContent(message)) {
+      handleTextMessage(message);
+      return;
+    }
+
+    if (messageUtils.isErrorMessageContent(message)) {
+      handleErrorMessage(message);
+      return;
+    }
+
+    // Handle tool call messages
+    if (messageUtils.isToolCallRequestEvent(message) || 
+        messageUtils.isToolCallExecutionEvent(message) || 
+        messageUtils.isToolCallSummaryMessage(message)) {
+      handleToolCallMessage(message);
+      return;
+    }
+
+    // Handle any other message types
+    handleOtherMessage(message);
+  };
+
+  return {
+    handleMessageEvent
+  };
+}; 

ui/src/lib/utils.ts

@@ -2,7 +2,7 @@ import { LLMCall } from "@/components/chat/LLMCallModal";
 import { ImageContent, TaskResultMessage } from "@/types/datamodel";
 import { clsx, type ClassValue } from "clsx";
 import { twMerge } from "tailwind-merge";
-import type { CompletionMessage, MemoryQueryEvent, ModelClientStreamingChunkEvent, TextMessageConfig, ToolCallExecutionEvent, ToolCallRequestEvent, ToolCallSummaryMessage } from "@/types/datamodel";
+import type { CompletionMessage, ErrorMessageConfig, MemoryQueryEvent, ModelClientStreamingChunkEvent, TextMessageConfig, ToolCallExecutionEvent, ToolCallRequestEvent, ToolCallSummaryMessage } from "@/types/datamodel";
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
@@ -184,6 +184,10 @@ export const messageUtils = {
     return typeof content === "object" && content !== null && "content" in content && "type" in content && content.type === "TextMessage";
   },
 
+  isErrorMessageContent(content: unknown): content is ErrorMessageConfig {
+    return typeof content === "object" && content !== null && "type" in content && content.type === "error";
+  },
+
   isUserTextMessageContent(content: unknown): content is TextMessageConfig {
     return messageUtils.isTextMessageContent(content) && content.source === "user";
   },

ui/src/types/datamodel.ts

@@ -70,6 +70,15 @@ export interface TextMessageConfig extends BaseMessageConfig {
   content: string;
 }
 
+export interface ErrorMessageConfig extends BaseMessageConfig {
+  data: {
+    task_result: TaskResult;
+    usage: string;
+    duration: number;
+  };
+  type: "error";
+}
+
 export interface MultiModalMessageConfig extends BaseMessageConfig {
   content: (string | ImageContent)[];
 }