Working example of streaming, intersection observer, other UI updates

src/interface/web/app/chat/chat.module.css

@@ -1,6 +1,6 @@
 div.main {
     height: 100vh;
-    color: black;
+    color: hsla(var(--foreground));
 }
 
 .suggestions {
@@ -58,7 +58,7 @@ div.chatBody {
 }
 
 .inputBox {
-    color: black;
+    color: hsla(var(--foreground));
 }
 
 div.chatLayout {

src/interface/web/app/chat/page.tsx

@@ -1,7 +1,7 @@
 'use client'
 
 import styles from './chat.module.css';
-import React, { Suspense, useEffect, useState } from 'react';
+import React, { Suspense, useEffect, useRef, useState } from 'react';
 
 import SuggestionCard from '../components/suggestions/suggestionCard';
 import SidePanel from '../components/sidePanel/chatHistorySidePanel';
@@ -10,7 +10,7 @@ import NavMenu from '../components/navMenu/navMenu';
 import { useSearchParams } from 'next/navigation'
 import Loading from '../components/loading/loading';
 
-import { setupWebSocket } from '../common/chatFunctions';
+import { handleCompiledReferences, handleImageResponse, setupWebSocket } from '../common/chatFunctions';
 
 import 'katex/dist/katex.min.css';
 import { Lightbulb, ArrowCircleUp, FileArrowUp, Microphone } from '@phosphor-icons/react';
@@ -18,13 +18,65 @@ import { Lightbulb, ArrowCircleUp, FileArrowUp, Microphone } from '@phosphor-ico
 import { Label } from "@/components/ui/label"
 import { Textarea } from "@/components/ui/textarea"
 import { Button } from '@/components/ui/button';
+import { Context, OnlineContextData, StreamMessage } from '../components/chatMessage/chatMessage';
+
+interface ChatInputProps {
+    sendMessage: (message: string) => void;
+    sendDisabled: boolean;
+}
+
+function ChatInputArea(props: ChatInputProps) {
+    const [message, setMessage] = useState('');
+
+    useEffect(() => {
+        if (message.startsWith('/')) {
+            const command = message.split(' ')[0].substring(1);
+            console.log('Command is:', command);
+        }
+    }, [message]);
+
+    function onSendMessage() {
+        props.sendMessage(message);
+        setMessage('');
+    }
 
-function TextareaWithLabel() {
     return (
-        <div className="grid w-full gap-1.5">
+        <>
-            {/* <Label htmlFor="message">Your message</Label> */}
+            <Button
-            <Textarea className='border-none min-h-[60px]' placeholder="Type / to see a list of commands" id="message" />
+                variant={'ghost'}
-        </div>
+                className="!bg-none p-1 h-auto text-3xl rounded-full text-gray-300 hover:text-gray-500"
+                disabled={props.sendDisabled}>
+                <FileArrowUp weight='fill' />
+            </Button>
+            <div className="grid w-full gap-1.5">
+                {/* <Label htmlFor="message">Your message</Label> */}
+                <Textarea
+                    className='border-none min-h-[60px]'
+                    placeholder="Type / to see a list of commands"
+                    id="message"
+                    value={message}
+                    onKeyDown={(e) => {
+                        if (e.key === 'Enter' && !e.shiftKey) {
+                            e.preventDefault();
+                            onSendMessage();
+                        }
+                    }}
+                    onChange={(e) => setMessage(e.target.value)}
+                    disabled={props.sendDisabled} />
+            </div>
+            <Button
+                variant={'ghost'}
+                className="!bg-none p-1 h-auto text-3xl rounded-full text-gray-300 hover:text-gray-500"
+                disabled={props.sendDisabled}>
+                <Microphone weight='fill' />
+            </Button>
+            <Button
+                className="bg-orange-300 hover:bg-orange-500 rounded-full p-0 h-auto text-3xl transition transform hover:-translate-y-1"
+                onClick={onSendMessage}
+                disabled={props.sendDisabled}>
+                <ArrowCircleUp />
+            </Button>
+        </>
     )
 }
 
@@ -37,12 +89,16 @@ interface ChatBodyDataProps {
     chatOptionsData: ChatOptions | null;
     setTitle: (title: string) => void;
     onConversationIdChange?: (conversationId: string) => void;
+    setQueryToProcess: (query: string) => void;
+    streamedMessages: StreamMessage[];
 }
 
 
 function ChatBodyData(props: ChatBodyDataProps) {
     const searchParams = useSearchParams();
     const conversationId = searchParams.get('conversationId');
+    const [message, setMessage] = useState('');
+    const [processingMessage, setProcessingMessage] = useState(false);
 
     useEffect(() => {
         if (conversationId) {
@@ -50,6 +106,34 @@ function ChatBodyData(props: ChatBodyDataProps) {
         }
     }, [conversationId, props.onConversationIdChange]);
 
+    // useEffect(() => {
+    //     // Reset the processing message whenever the streamed messages are updated
+    //     if (props.streamedMessages) {
+    //         setProcessingMessage(false);
+    //     }
+    // }, [props.streamedMessages]);
+
+    useEffect(() => {
+        if (message) {
+            setProcessingMessage(true);
+            props.setQueryToProcess(message);
+            // setTimeout(() => {
+            //     setProcessingMessage(false);
+            // }, 1000);
+        }
+    }, [message]);
+
+    useEffect(() => {
+        if (props.streamedMessages &&
+            props.streamedMessages.length > 0 &&
+            props.streamedMessages[props.streamedMessages.length - 1].completed) {
+
+            setProcessingMessage(false);
+        } else {
+            setMessage('');
+        }
+    }, [props.streamedMessages]);
+
     if (!conversationId) {
         return (
             <div className={styles.suggestions}>
@@ -69,37 +153,91 @@ function ChatBodyData(props: ChatBodyDataProps) {
     return (
         <>
             <div className={false ? styles.chatBody : styles.chatBodyFull}>
-                <ChatHistory conversationId={conversationId} setTitle={props.setTitle} />
+                <ChatHistory conversationId={conversationId} setTitle={props.setTitle} pendingMessage={processingMessage ? message : ''} incomingMessages={props.streamedMessages} />
             </div>
             <div className={`${styles.inputBox} bg-background align-middle items-center justify-center`}>
-                <Button variant={'ghost'} className="!bg-none p-1 h-auto text-3xl rounded-full !hover:bg-accent ">
+                <ChatInputArea sendMessage={(message) => setMessage(message)} sendDisabled={processingMessage} />
-                    <FileArrowUp fill="hsl(var(--accent-foreground))" />
-                </Button>
-                <TextareaWithLabel />
-                <Button variant={'ghost'} className="!bg-none p-1 h-auto text-3xl rounded-full !hover:bg-accent">
-                    <Microphone fill="hsl(var(--accent-foreground))" />
-                </Button>
-                <Button className="bg-orange-300 hover:bg-orange-500 rounded-full p-0 h-auto text-3xl">
-                    <ArrowCircleUp />
-                </Button>
-                {/* <input className={styles.inputBox} type="text" placeholder="Type / to see a list of commands" onInput={(e) => handleChatInput(e)} /> */}
-                {/* <button className={styles.inputBox}>Send</button> */}
             </div>
         </>
     );
 }
 
-function handleChatInput(e: React.FormEvent<HTMLInputElement>) {
-    const target = e.target as HTMLInputElement;
-    console.log(target.value);
-}
-
 export default function Chat() {
     const [chatOptionsData, setChatOptionsData] = useState<ChatOptions | null>(null);
     const [isLoading, setLoading] = useState(true);
-    const [title, setTitle] = useState('Chat');
+    const [title, setTitle] = useState('Khoj AI - Chat');
     const [conversationId, setConversationID] = useState<string | null>(null);
     const [chatWS, setChatWS] = useState<WebSocket | null>(null);
+    const [messages, setMessages] = useState<StreamMessage[]>([]);
+    const [queryToProcess, setQueryToProcess] = useState<string>('');
+    const [processQuerySignal, setProcessQuerySignal] = useState(false);
+
+
+    const handleWebSocketMessage = (event: MessageEvent) => {
+        let chunk = event.data;
+
+        let currentMessage = messages.find(message => !message.completed);
+
+        if (!currentMessage) {
+            console.error("No current message found");
+            return;
+        }
+
+        // Process WebSocket streamed data
+        if (chunk === "start_llm_response") {
+            console.log("Started streaming", new Date());
+        } else if (chunk === "end_llm_response") {
+            currentMessage.completed = true;
+        } else {
+            // Get the current message
+            // Process and update state with the new message
+            if (chunk.includes("application/json")) {
+                chunk = JSON.parse(chunk);
+            }
+
+            const contentType = chunk["content-type"];
+            if (contentType === "application/json") {
+                try {
+                    if (chunk.image && chunk.detail) {
+                        let responseWithReference = handleImageResponse(chunk);
+                        console.log("Image response", responseWithReference);
+                        if (responseWithReference.response) currentMessage.rawResponse = responseWithReference.response;
+                        if (responseWithReference.online) currentMessage.onlineContext = responseWithReference.online;
+                        if (responseWithReference.context) currentMessage.context = responseWithReference.context;
+                    } else if (chunk.type == "status") {
+                        currentMessage.trainOfThought.push(chunk.message);
+                    } else if (chunk.type == "rate_limit") {
+                        console.log("Rate limit message", chunk);
+                        currentMessage.rawResponse = chunk.message;
+                    } else {
+                        console.log("any message", chunk);
+                    }
+                } catch (error) {
+                    console.error("Error processing message", error);
+                    currentMessage.completed = true;
+                } finally {
+                    // no-op
+                }
+
+            } else {
+                // Update the current message with the new chunk
+                if (chunk && chunk.includes("### compiled references:")) {
+                    let responseWithReference = handleCompiledReferences(chunk, "");
+                    currentMessage.rawResponse += responseWithReference.response;
+
+                    if (responseWithReference.response) currentMessage.rawResponse = responseWithReference.response;
+                    if (responseWithReference.online) currentMessage.onlineContext = responseWithReference.online;
+                    if (responseWithReference.context) currentMessage.context = responseWithReference.context;
+                } else {
+                    // If the chunk is not a JSON object, just display it as is
+                    currentMessage.rawResponse += chunk;
+                }
+
+            }
+        };
+        // Update the state with the new message, currentMessage
+        setMessages([...messages]);
+    }
 
     useEffect(() => {
         fetch('/api/chat/options')
@@ -108,7 +246,6 @@ export default function Chat() {
                 setLoading(false);
                 // Render chat options, if any
                 if (data) {
-                    console.log(data);
                     setChatOptionsData(data);
                 }
             })
@@ -119,6 +256,43 @@ export default function Chat() {
 
     }, []);
 
+    useEffect(() => {
+        if (chatWS && queryToProcess) {
+            // Add a new object to the state
+            const newStreamMessage: StreamMessage = {
+                rawResponse: "",
+                trainOfThought: [],
+                context: [],
+                onlineContext: {},
+                completed: false,
+                timestamp: (new Date()).toISOString(),
+                rawQuery: queryToProcess || "",
+            }
+            setMessages(prevMessages => [...prevMessages, newStreamMessage]);
+            console.log("Messages", messages);
+            setProcessQuerySignal(true);
+        } else {
+            if (!chatWS) {
+                console.error("No WebSocket connection available");
+            }
+            if (!queryToProcess) {
+                console.error("No query to process");
+            }
+        }
+    }, [queryToProcess]);
+
+    // useEffect(() => {
+    //     console.log("messages", messages);
+    // }, [messages]);
+
+    useEffect(() => {
+        if (processQuerySignal && chatWS) {
+            setProcessQuerySignal(false);
+            chatWS.onmessage = handleWebSocketMessage;
+            chatWS?.send(queryToProcess);
+        }
+    }, [processQuerySignal]);
+
     useEffect(() => {
         (async () => {
             if (conversationId) {
@@ -141,22 +315,24 @@ export default function Chat() {
     return (
         <div className={styles.main + " " + styles.chatLayout}>
             <title>
-                Khoj AI - {title}
+                {title}
             </title>
-                <div className={styles.sidePanel}>
+            <div className={styles.sidePanel}>
-                    <SidePanel webSocketConnected={chatWS !== null} conversationId={conversationId} />
+                <SidePanel webSocketConnected={chatWS !== null} conversationId={conversationId} />
-                </div>
+            </div>
-                <div className={styles.chatBox}>
+            <div className={styles.chatBox}>
-                    <NavMenu selected="Chat" title={title} />
+                <NavMenu selected="Chat" title={title} />
-                    <div className={styles.chatBoxBody}>
+                <div className={styles.chatBoxBody}>
-                        <Suspense fallback={<Loading />}>
+                    <Suspense fallback={<Loading />}>
-                            <ChatBodyData
+                        <ChatBodyData
-                                chatOptionsData={chatOptionsData}
+                            streamedMessages={messages}
-                                setTitle={setTitle}
+                            chatOptionsData={chatOptionsData}
-                                onConversationIdChange={handleConversationIdChange} />
+                            setTitle={setTitle}
-                        </Suspense>
+                            setQueryToProcess={setQueryToProcess}
-                    </div>
+                            onConversationIdChange={handleConversationIdChange} />
+                    </Suspense>
                 </div>
+            </div>
         </div>
     )
 }

src/interface/web/app/common/chatFunctions.ts

@@ -8,7 +8,7 @@ interface ResponseWithReferences {
     response?: string;
 }
 
-function handleCompiledReferences(chunk: string, currentResponse: string) {
+export function handleCompiledReferences(chunk: string, currentResponse: string) {
     const rawReference = chunk.split("### compiled references:")[1];
     const rawResponse = chunk.split("### compiled references:")[0];
     let references: ResponseWithReferences = {};
@@ -69,10 +69,6 @@ async function sendChatStream(
         }
 }
 
-export function sendChatWS(websocket: WebSocket, message: string) {
-    websocket.send(message);
-}
-
 export const setupWebSocket = async (conversationId: string) => {
     const wsProtocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
 
@@ -108,3 +104,42 @@ export const setupWebSocket = async (conversationId: string) => {
 
     return chatWS;
 };
+
+export function handleImageResponse(imageJson: any) {
+
+    let rawResponse = "";
+
+    if (imageJson.image) {
+        const inferredQuery = imageJson.inferredQueries?.[0] ?? "generated image";
+
+        // If response has image field, response is a generated image.
+        if (imageJson.intentType === "text-to-image") {
+            rawResponse += `![generated_image](data:image/png;base64,${imageJson.image})`;
+        } else if (imageJson.intentType === "text-to-image2") {
+            rawResponse += `![generated_image](${imageJson.image})`;
+        } else if (imageJson.intentType === "text-to-image-v3") {
+            rawResponse = `![](data:image/webp;base64,${imageJson.image})`;
+        }
+        if (inferredQuery) {
+            rawResponse += `\n\n**Inferred Query**:\n\n${inferredQuery}`;
+        }
+    }
+
+    let reference: ResponseWithReferences = {};
+
+    if (imageJson.context && imageJson.context.length > 0) {
+        const rawReferenceAsJson = imageJson.context;
+        if (rawReferenceAsJson instanceof Array) {
+            reference.context = rawReferenceAsJson;
+        } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
+            reference.online = rawReferenceAsJson;
+        }
+    }
+    if (imageJson.detail) {
+        // The detail field contains the improved image prompt
+        rawResponse += imageJson.detail;
+    }
+
+    reference.response = rawResponse;
+    return reference;
+}

src/interface/web/app/components/chatHistory/chatHistory.module.css

@@ -20,3 +20,11 @@ div.agentIndicator a {
 div.agentIndicator {
     padding: 10px;
 }
+
+div.trainOfThought {
+    border: 1px var(--border-color) solid;
+    border-radius: 16px;
+    padding: 16px;
+    margin: 12px;
+    box-shadow: 0 4px 10px rgba(0, 0, 0, 0.03);
+}

src/interface/web/app/components/chatHistory/chatHistory.tsx

@@ -3,9 +3,9 @@
 import styles from './chatHistory.module.css';
 import { useRef, useEffect, useState } from 'react';
 
-import ChatMessage, { ChatHistoryData, SingleChatMessage } from '../chatMessage/chatMessage';
+import ChatMessage, { ChatHistoryData, SingleChatMessage, StreamMessage, TrainOfThought } from '../chatMessage/chatMessage';
 
-import ReferencePanel, { hasValidReferences} from '../referencePanel/referencePanel';
+import ReferencePanel, { hasValidReferences } from '../referencePanel/referencePanel';
 
 import { ScrollArea } from "@/components/ui/scroll-area"
 
@@ -29,43 +29,149 @@ interface ChatHistory {
 interface ChatHistoryProps {
     conversationId: string;
     setTitle: (title: string) => void;
+    incomingMessages?: StreamMessage[];
+    pendingMessage?: string;
+}
+
+
+function constructTrainOfThought(trainOfThought: string[], lastMessage: boolean, key: string) {
+    const lastIndex = trainOfThought.length - 1;
+    return (
+        <div className={`${styles.trainOfThought}`} key={key}>
+            {trainOfThought.map((train, index) => (
+                <TrainOfThought message={train} primary={index === lastIndex && lastMessage} />
+            ))}
+        </div>
+    )
 }
 
 
 export default function ChatHistory(props: ChatHistoryProps) {
     const [data, setData] = useState<ChatHistoryData | null>(null);
-    const [isLoading, setLoading] = useState(true)
+    const [isLoading, setLoading] = useState(true);
+    const [currentPage, setCurrentPage] = useState(0);
+    const [hasMoreMessages, setHasMoreMessages] = useState(true);
+
     const ref = useRef<HTMLDivElement>(null);
-    const chatHistoryRef = useRef(null);
+    const chatHistoryRef = useRef<HTMLDivElement | null>(null);
+    const sentinelRef = useRef<HTMLDivElement | null>(null);
 
     const [showReferencePanel, setShowReferencePanel] = useState(true);
     const [referencePanelData, setReferencePanelData] = useState<SingleChatMessage | null>(null);
+    const [incompleteIncomingMessageIndex, setIncompleteIncomingMessageIndex] = useState<number | null>(null);
 
-	useEffect(() => {
+    // useEffect(() => {
 
-        fetch(`/api/chat/history?client=web&conversation_id=${props.conversationId}&n=10`)
+    //     // TODO add intersection observer to load more messages incrementally using parameter n=. Right now, it loads all messages at once.
+
+    //     fetch(`/api/chat/history?client=web&conversation_id=${props.conversationId}`)
+    //         .then(response => response.json())
+    //         .then((chatData: ChatResponse) => {
+    //             setLoading(false);
+
+    //             // Render chat options, if any
+    //             if (chatData) {
+    //                 setData(chatData.response);
+    //                 props.setTitle(chatData.response.slug);
+    //             }
+    //         })
+    //         .catch(err => {
+    //             console.error(err);
+    //             return;
+    //         });
+    // }, [props.conversationId]);
+
+    useEffect(() => {
+        console.log("hasMoreMessages", hasMoreMessages);
+        const observer = new IntersectionObserver(entries => {
+            console.log("entries intersection observer", entries);
+            if (entries[0].isIntersecting && hasMoreMessages) {
+                console.log("call fetchMoreMessages");
+                fetchMoreMessages(currentPage);
+                console.log("currentPage", currentPage);
+                setCurrentPage((prev) => prev + 1);
+            }
+        }, { threshold: 1.0 });
+
+        if (sentinelRef.current) {
+            console.log("observe sentinel");
+            observer.observe(sentinelRef.current);
+        }
+
+        return () => observer.disconnect();
+    }, [sentinelRef.current, hasMoreMessages, currentPage, props.conversationId]);
+
+    const fetchMoreMessages = (currentPage: number) => {
+        if (!hasMoreMessages) return;
+
+        console.log("fetchMoreMessages", currentPage);
+
+        const nextPage = currentPage + 1;
+        fetch(`/api/chat/history?client=web&conversation_id=${props.conversationId}&n=${10*nextPage}`)
             .then(response => response.json())
             .then((chatData: ChatResponse) => {
-                setLoading(false);
+                console.log(chatData);
-
+                if (chatData && chatData.response && chatData.response.chat.length > 0) {
-                // Render chat options, if any
-                if (chatData) {
                     console.log(chatData);
+
+                    if (chatData.response.chat.length === data?.chat.length) {
+                        setHasMoreMessages(false);
+                        return;
+                    }
+
+                    scrollToBottom();
+
                     setData(chatData.response);
-                    props.setTitle(chatData.response.slug);
+                    setLoading(false);
+                } else {
+                    console.log("No more messages");
+                    setHasMoreMessages(false);
                 }
             })
             .catch(err => {
                 console.error(err);
-                return;
             });
-	}, [props.conversationId]);
+    };
+
+    useEffect(() => {
+        if (props.incomingMessages) {
+            const lastMessage = props.incomingMessages[props.incomingMessages.length - 1];
+            if (lastMessage && !lastMessage.completed) {
+                setIncompleteIncomingMessageIndex(props.incomingMessages.length - 1);
+            }
+        }
+
+        console.log("isUserAtBottom", isUserAtBottom());
+
+        if (isUserAtBottom()) {
+            scrollToBottom();
+        }
+
+    }, [props.incomingMessages]);
+
+    const scrollToBottom = () => {
+        if (chatHistoryRef.current) {
+            chatHistoryRef.current.scrollIntoView(false);
+        }
+    }
+
+    const isUserAtBottom = () => {
+        if (!chatHistoryRef.current) return false;
+
+        // NOTE: This isn't working. It always seems to return true. This is because
+
+        const { scrollTop, scrollHeight, clientHeight } = chatHistoryRef.current as HTMLDivElement;
+        const threshold = 25; // pixels from the bottom
+
+        // Considered at the bottom if within threshold pixels from the bottom
+        return scrollTop + clientHeight >= scrollHeight - threshold;
+    }
 
 
     useEffect(() => {
         const observer = new MutationObserver((mutationsList, observer) => {
             // If the addedNodes property has one or more nodes
-            for(let mutation of mutationsList) {
+            for (let mutation of mutationsList) {
                 if (mutation.type === 'childList' && mutation.addedNodes.length > 0) {
                     // Call your function here
                     renderMathInElement(document.body, {
@@ -88,9 +194,9 @@ export default function ChatHistory(props: ChatHistoryProps) {
         return () => observer.disconnect();
     }, []);
 
-    if (isLoading) {
+    // if (isLoading) {
-        return <Loading />;
+    //     return <Loading />;
-    }
+    // }
 
     function constructAgentLink() {
         if (!data || !data.agent || !data.agent.slug) return `/agents`;
@@ -111,9 +217,10 @@ export default function ChatHistory(props: ChatHistoryProps) {
         <ScrollArea className={`h-[80vh]`}>
             <div ref={ref}>
                 <div className={styles.chatHistory} ref={chatHistoryRef}>
+                    <div ref={sentinelRef} style={{ height: '1px' }}></div>
                     {(data && data.chat) && data.chat.map((chatMessage, index) => (
                         <ChatMessage
-                            key={index}
+                            key={`${index}fullHistory`}
                             chatMessage={chatMessage}
                             setReferencePanelData={setReferencePanelData}
                             setShowReferencePanel={setShowReferencePanel}
@@ -121,9 +228,77 @@ export default function ChatHistory(props: ChatHistoryProps) {
                             borderLeftColor='orange-500'
                         />
                     ))}
+                    {
+                        props.incomingMessages && props.incomingMessages.map((message, index) => {
+                            return (
+                                <>
+                                    <ChatMessage
+                                        key={`${index}outgoing`}
+                                        chatMessage={
+                                            {
+                                                message: message.rawQuery,
+                                                context: [],
+                                                onlineContext: {},
+                                                created: message.timestamp,
+                                                by: "you",
+                                                intent: {},
+                                                automationId: '',
+
+                                            }
+                                        }
+                                        setReferencePanelData={() => { }}
+                                        setShowReferencePanel={() => { }}
+                                        customClassName='fullHistory'
+                                        borderLeftColor='orange-500' />
+                                    {
+                                        message.trainOfThought && constructTrainOfThought(message.trainOfThought, index === incompleteIncomingMessageIndex, `${index}trainOfThought`)
+                                    }
+                                    <ChatMessage
+                                        key={`${index}incoming`}
+                                        chatMessage={
+                                            {
+                                                message: message.rawResponse,
+                                                context: message.context,
+                                                onlineContext: message.onlineContext,
+                                                created: message.timestamp,
+                                                by: "khoj",
+                                                intent: {},
+                                                automationId: '',
+                                            }
+                                        }
+                                        setReferencePanelData={setReferencePanelData}
+                                        setShowReferencePanel={setShowReferencePanel}
+                                        customClassName='fullHistory'
+                                        borderLeftColor='orange-500'
+                                    />
+                                </>
+                            )
+                        })
+                    }
+                    {
+                        props.pendingMessage &&
+                        <ChatMessage
+                            key={"pendingMessage"}
+                            chatMessage={
+                                {
+                                    message: props.pendingMessage,
+                                    context: [],
+                                    onlineContext: {},
+                                    created: new Date().toISOString(),
+                                    by: "you",
+                                    intent: {},
+                                    automationId: '',
+                                }
+                            }
+                            setReferencePanelData={() => { }}
+                            setShowReferencePanel={() => { }}
+                            customClassName='fullHistory'
+                            borderLeftColor='orange-500'
+                        />
+                    }
                     {
                         (hasValidReferences(referencePanelData) && showReferencePanel) &&
-                            <ReferencePanel referencePanelData={referencePanelData} setShowReferencePanel={setShowReferencePanel} />
+                        <ReferencePanel referencePanelData={referencePanelData} setShowReferencePanel={setShowReferencePanel} />
                     }
                     <div className={`${styles.agentIndicator}`}>
                         <a className='no-underline mx-2 flex' href={constructAgentLink()} target="_blank" rel="noreferrer">

src/interface/web/app/components/chatMessage/chatMessage.module.css

@@ -110,6 +110,19 @@ button.copyButton img {
     width: 24px;
 }
 
+div.trainOfThought strong {
+    font-weight: 500;
+}
+
+div.trainOfThought.primary strong {
+    font-weight: 500;
+    color: hsla(var(--secondary-foreground));
+}
+
+div.trainOfThought.primary p {
+    color: inherit;
+}
+
 @media screen and (max-width: 768px) {
     div.youfullHistory {
         max-width: 100%;

src/interface/web/app/components/chatMessage/chatMessage.tsx

@@ -12,7 +12,7 @@ import 'highlight.js/styles/github.css'
 
 import { hasValidReferences } from '../referencePanel/referencePanel';
 
-import { ThumbsUp, ThumbsDown, Copy } from '@phosphor-icons/react';
+import { ThumbsUp, ThumbsDown, Copy, Brain, Cloud, Folder, Book } from '@phosphor-icons/react';
 
 const md = new markdownIt({
     html: true,
@@ -96,6 +96,19 @@ export interface SingleChatMessage {
     }
 }
 
+export interface StreamMessage {
+    rawResponse: string;
+    trainOfThought: string[];
+    context: Context[];
+    onlineContext: {
+        [key: string]: OnlineContextData
+    }
+    completed: boolean;
+    rawQuery: string;
+    timestamp: string;
+}
+
+
 export interface ChatHistoryData {
     chat: SingleChatMessage[];
     agent: AgentData;
@@ -117,7 +130,7 @@ function FeedbackButtons() {
 }
 
 function onClickMessage(event: React.MouseEvent<any>, chatMessage: SingleChatMessage, setReferencePanelData: Function, setShowReferencePanel: Function) {
-    console.log("Clicked on message", chatMessage);
+    // console.log("Clicked on message", chatMessage);
     setReferencePanelData(chatMessage);
     setShowReferencePanel(true);
 }
@@ -130,6 +143,51 @@ interface ChatMessageProps {
     borderLeftColor?: string;
 }
 
+interface TrainOfThoughtProps {
+    message: string;
+    primary: boolean;
+}
+
+function chooseIconFromHeader(header: string, iconColor: string) {
+    const compareHeader = header.toLowerCase();
+    if (compareHeader.includes("understanding")) {
+        return <Brain className={`inline mr-2 ${iconColor}`} />
+    }
+
+    if (compareHeader.includes("generating")) {
+        return <Cloud className={`inline mr-2 ${iconColor}`} />;
+    }
+
+    if (compareHeader.includes("data sources")) {
+        return <Folder className={`inline mr-2 ${iconColor}`} />;
+    }
+
+    if (compareHeader.includes("notes")) {
+        return <Folder className={`inline mr-2 ${iconColor}`} />;
+    }
+
+    if (compareHeader.includes("read")) {
+        return <Book className={`inline mr-2 ${iconColor}`} />;
+    }
+
+    return <Brain className={`inline mr-2 ${iconColor}`} />;
+}
+
+export function TrainOfThought(props: TrainOfThoughtProps) {
+    // The train of thought comes in as a markdown-formatted string. It starts with a heading delimited by two asterisks at the start and end and a colon, followed by the message. Example: **header**: status. This function will parse the message and render it as a div.
+    let extractedHeader = props.message.match(/\*\*(.*)\*\*/);
+    let header = extractedHeader ? extractedHeader[1] : "";
+    const iconColor = props.primary ? 'text-orange-400' : 'text-gray-500';
+    const icon = chooseIconFromHeader(header, iconColor);
+    let markdownRendered = md.render(props.message);
+    return (
+        <div className={`flex items-center ${props.primary ? 'text-gray-400' : 'text-gray-300'} ${styles.trainOfThought} ${props.primary ? styles.primary : ''}`} >
+            {icon}
+            <div dangerouslySetInnerHTML={{ __html: markdownRendered }} />
+        </div>
+    )
+}
+
 export default function ChatMessage(props: ChatMessageProps) {
     const [copySuccess, setCopySuccess] = useState<boolean>(false);
 
@@ -154,7 +212,6 @@ export default function ChatMessage(props: ChatMessageProps) {
     useEffect(() => {
         if (messageRef.current) {
             const preElements = messageRef.current.querySelectorAll('pre > .hljs');
-            console.log("make copy button");
             preElements.forEach((preElement) => {
                 const copyButton = document.createElement('button');
                 const copyImage = document.createElement('img');

src/interface/web/app/components/sidePanel/chatHistorySidePanel.tsx

@@ -77,7 +77,6 @@ function renameConversation(conversationId: string, newTitle: string) {
     })
         .then(response => response.json())
         .then(data => {
-            console.log(data);
         })
         .catch(err => {
             console.error(err);
@@ -96,7 +95,6 @@ function shareConversation(conversationId: string, setShareUrl: (url: string) =>
     })
         .then(response => response.json())
         .then(data => {
-            console.log(data);
             setShareUrl(data.url);
         })
         .catch(err => {
@@ -116,7 +114,6 @@ function deleteConversation(conversationId: string) {
     })
         .then(response => response.json())
         .then(data => {
-            console.log(data);
         })
         .catch(err => {
             console.error(err);
@@ -153,7 +150,6 @@ function modifyFileFilterForConversation(
         .then(response => response.json())
         .then(data => {
             setAddedFiles(data);
-            console.log(data);
         })
         .catch(err => {
             console.error(err);
@@ -417,7 +413,6 @@ function ChatSessionActionMenu(props: ChatSessionActionMenuProps) {
                             <Button
                                 onClick={() => {
                                     navigator.clipboard.writeText(shareUrl);
-                                    console.log("shared");
                                 }}
                                 variant={'default'}>Copy</Button>
                         }
@@ -428,7 +423,6 @@ function ChatSessionActionMenu(props: ChatSessionActionMenuProps) {
     }
 
     if (isDeleting) {
-        console.log("Deleting");
         return (
             <AlertDialog
                 open={isDeleting}