Better layout for terminal and chat messages

surfsense_backend/app/services/streaming_service.py

@@ -43,17 +43,7 @@ class StreamingService:
         self.message_annotations[0]["content"].append(message)
 
         # Return only the delta annotation
-        # annotation = {"type": "TERMINAL_INFO", "content": [message]}
+        annotation = {"type": "TERMINAL_INFO", "data": message}
-        agent_data = {
-            "agent_name": "Terminal agent",
-            "events": [
-                {
-                    "status": message_type,
-                    "result": message.get("text", ""),
-                }
-            ],
-        }
-        annotation = {"type": "agent_events", "data": agent_data}
         return f"8:[{json.dumps(annotation)}]\n"
 
     def format_sources_delta(self, sources: List[Dict[str, Any]]) -> str:

surfsense_web/components/chat_v2/ChatInterface.tsx

@@ -7,6 +7,7 @@ import {
     ChatMessages,
     useChatUI,
     ChatMessage,
+    Message,
 } from "@llamaindex/chat-ui";
 import { Document } from "@/hooks/use-documents";
 import { CustomChatInput } from "@/components/chat_v2/ChatInputGroup";
@@ -26,23 +27,48 @@ interface ChatInterfaceProps {
     onResearchModeChange?: (mode: ResearchMode) => void;
 }
 
-function ChatMessageDisplay() {
+function ChatMessageDisplay({
+    message,
+    isLast,
+}: {
+    message: Message;
+    isLast: boolean;
+}) {
+    return (
+        <ChatMessage
+            message={message}
+            isLast={isLast}
+            className="flex flex-col "
+        >
+            {message.role === "assistant" ? (
+                <div className="flex-1 flex flex-col space-y-4">
+                    <TerminalDisplay message={message} />
+                    <ChatMessage.Content className="flex-1">
+                        <ChatMessage.Content.Markdown />
+                    </ChatMessage.Content>
+                    <ChatMessage.Actions className="flex-1 flex flex-row justify-end" />
+                </div>
+            ) : (
+                <ChatMessage.Content className="flex-1">
+                    <ChatMessage.Content.Markdown />
+                </ChatMessage.Content>
+            )}
+        </ChatMessage>
+    );
+}
+
+function ChatMessagesDisplay() {
     const { messages } = useChatUI();
 
     return (
         <ChatMessages className="flex-1">
             <ChatMessages.List className="p-4">
                 {messages.map((message, index) => (
-                    <div key={`Message-${index}`}>
+                    <ChatMessageDisplay
-                        {message.role === "assistant" && (
+                        key={`Message-${index}`}
-                            <TerminalDisplay messages={messages} />
+                        message={message}
-                        )}
+                        isLast={index === messages.length - 1}
-                        <ChatMessage
+                    />
-                            key={index}
-                            message={message}
-                            isLast={index === messages.length - 1}
-                        />
-                    </div>
                 ))}
             </ChatMessages.List>
             <ChatMessages.Loading />
@@ -64,7 +90,7 @@ export default function ChatInterface({
     return (
         <ChatSection handler={handler} className="flex h-full">
             <div className="flex flex-1 flex-col">
-                <ChatMessageDisplay />
+                <ChatMessagesDisplay />
                 <div className="border-t p-4">
                     <CustomChatInput
                         onDocumentSelectionChange={onDocumentSelectionChange}

surfsense_web/components/chat_v2/ChatTerminal.tsx

@@ -3,31 +3,24 @@
 import React from "react";
 import { getAnnotationData, Message } from "@llamaindex/chat-ui";
 
-export default function TerminalDisplay({ messages }: { messages: Message[] }) {
+export default function TerminalDisplay({ message }: { message: Message }) {
     const [isCollapsed, setIsCollapsed] = React.useState(true);
 
     // Get the last assistant message that's not being typed
-    const lastAssistantMessage = messages
+    if (!message) {
-        .filter((msg) => msg.role === "assistant")
-        .pop();
-
-    if (!lastAssistantMessage) {
         return <></>;
     }
 
-    const annotations = getAnnotationData(lastAssistantMessage, "agent_events");
+    interface TerminalInfo {
-
+        id: number;
-    if (annotations.length === 0) {
+        text: string;
-        return <></>;
+        type: string;
     }
 
-    // Get all agent events and combine them
+    const events = getAnnotationData(
-    const events: any[] = [];
+        message,
-    annotations.forEach((annotation: any) => {
+        "TERMINAL_INFO"
-        if (annotation?.events && Array.isArray(annotation.events)) {
+    ) as TerminalInfo[];
-            events.push(...annotation.events);
-        }
-    });
 
     if (events.length === 0) {
         return <></>;
@@ -85,16 +78,17 @@ export default function TerminalDisplay({ messages }: { messages: Message[] }) {
             {!isCollapsed && (
                 <div className="h-64 overflow-y-auto p-4 space-y-1 bg-gray-900">
                     {events.map((event, index) => (
-                        <div key={index} className="text-green-400">
+                        <div
+                            key={`${event.id}-${index}`}
+                            className="text-green-400"
+                        >
                             <span className="text-blue-400">$</span>
                             <span className="text-yellow-400 ml-2">
-                                [{event.status || "completed"}]
+                                [{event.type || ""}]
+                            </span>
+                            <span className="text-gray-300 ml-4 mt-1 pl-2 border-l-2 border-gray-600">
+                                {event.text || ""}...
                             </span>
-                            {event.result && (
-                                <span className="text-gray-300 ml-4 mt-1 pl-2 border-l-2 border-gray-600">
-                                    {event.result.slice(0, 100)}...
-                                </span>
-                            )}
                         </div>
                     ))}
                     {events.length === 0 && (