feat(ChatWindow): Implement message buffering for improved UI updates and token handling

src/components/ChatWindow.tsx

@@ -420,6 +420,9 @@ const ChatWindow = ({ id }: { id?: string }) => {
 
     let sources: Document[] | undefined = undefined;
     let recievedMessage = '';
+    let messageBuffer = '';
+    let tokenCount = 0;
+    const bufferThreshold = 10;
     let added = false;
     let messageChatHistory = chatHistory;
 
@@ -542,11 +545,18 @@ const ChatWindow = ({ id }: { id?: string }) => {
       }
 
       if (data.type === 'response') {
+        // Add to buffer instead of immediately updating UI
+        messageBuffer += data.data;
+        recievedMessage += data.data;
+        tokenCount++;
+
+        // Only update UI every bufferThreshold tokens
+        if (tokenCount >= bufferThreshold) {
           if (!added) {
             setMessages((prevMessages) => [
               ...prevMessages,
               {
-              content: data.data,
+                content: messageBuffer,
                 messageId: data.messageId, // Use the AI message ID from the backend
                 chatId: chatId!,
                 role: 'assistant',
@@ -559,33 +569,31 @@ const ChatWindow = ({ id }: { id?: string }) => {
             setMessages((prev) =>
               prev.map((message) => {
                 if (message.messageId === data.messageId) {
-                return { ...message, content: message.content + data.data };
+                  return { ...message, content: recievedMessage };
                 }
                 return message;
               }),
             );
           }
 
-        recievedMessage += data.data;
+          // Reset buffer and counter
+          messageBuffer = '';
+          tokenCount = 0;
           setScrollTrigger((prev) => prev + 1);
         }
+      }
 
       if (data.type === 'messageEnd') {
         // Clear analysis progress
         setAnalysisProgress(null);
 
-        setChatHistory((prevHistory) => [
+        // Ensure final message content is displayed (flush any remaining buffer)
-          ...prevHistory,
-          ['human', message],
-          ['assistant', recievedMessage],
-        ]);
-
-        // Always update the message, adding modelStats if available
         setMessages((prev) =>
           prev.map((message) => {
             if (message.messageId === data.messageId) {
               return {
                 ...message,
+                content: recievedMessage, // Use the complete received message
                 // Include model stats if available, otherwise null
                 modelStats: data.modelStats || null,
                 // Make sure the searchQuery is preserved (if available in the message data)
@@ -597,6 +605,12 @@ const ChatWindow = ({ id }: { id?: string }) => {
           }),
         );
 
+        setChatHistory((prevHistory) => [
+          ...prevHistory,
+          ['human', message],
+          ['assistant', recievedMessage],
+        ]);
+
         setLoading(false);
         setScrollTrigger((prev) => prev + 1);
 

src/components/CitationLink.tsx

@@ -1,5 +1,6 @@
 import { Document } from '@langchain/core/documents';
-import { useState } from 'react';
+import { useState, useRef } from 'react';
+import { createPortal } from 'react-dom';
 import MessageSource from './MessageSource';
 
 interface CitationLinkProps {
@@ -10,6 +11,9 @@ interface CitationLinkProps {
 
 const CitationLink = ({ number, source, url }: CitationLinkProps) => {
   const [showTooltip, setShowTooltip] = useState(false);
+  const [tooltipPosition, setTooltipPosition] = useState({ x: 0, y: 0 });
+  const spanRef = useRef<HTMLSpanElement>(null);
+
   const linkContent = (
     <a
       href={url}
@@ -21,19 +25,45 @@ const CitationLink = ({ number, source, url }: CitationLinkProps) => {
     </a>
   );
 
+  const handleMouseEnter = () => {
+    if (spanRef.current) {
+      const rect = spanRef.current.getBoundingClientRect();
+      setTooltipPosition({
+        x: rect.left + rect.width / 2,
+        y: rect.top,
+      });
+      setShowTooltip(true);
+    }
+  };
+
+  const handleMouseLeave = () => {
+    setShowTooltip(false);
+  };
+
   // If we have source data, wrap with tooltip
   if (source) {
     return (
-      <div className="relative inline-block">
+      <>
-        <div
+        <span
-          onMouseEnter={() => setShowTooltip(true)}
+          ref={spanRef}
-          onMouseLeave={() => setShowTooltip(false)}
+          className="relative inline-block"
+          onMouseEnter={handleMouseEnter}
+          onMouseLeave={handleMouseLeave}
         >
           {linkContent}
-        </div>
+        </span>
 
-        {showTooltip && (
+        {showTooltip &&
-          <div className="absolute z-50 bottom-full mb-2 left-1/2 transform -translate-x-1/2 animate-in fade-in-0 duration-150">
+          typeof window !== 'undefined' &&
+          createPortal(
+            <div
+              className="fixed z-50 animate-in fade-in-0 duration-150"
+              style={{
+                left: tooltipPosition.x,
+                top: tooltipPosition.y - 8,
+                transform: 'translate(-50%, -100%)',
+              }}
+            >
               <div className="bg-light-primary dark:bg-dark-primary border rounded-lg border-light-200 dark:border-dark-200 shadow-lg w-96">
                 <MessageSource
                   source={source}
@@ -42,9 +72,10 @@ const CitationLink = ({ number, source, url }: CitationLinkProps) => {
               </div>
               {/* Tooltip arrow */}
               <div className="absolute top-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-l-4 border-r-4 border-t-4 border-transparent border-t-light-200 dark:border-t-dark-200"></div>
-          </div>
+            </div>,
+            document.body,
           )}
-      </div>
+      </>
     );
   }
 

src/components/MessageTabs.tsx

@@ -100,7 +100,6 @@ const MessageTabs = ({
 
   // Process message content
   useEffect(() => {
-    const citationRegex = /\[([^\]]+)\]/g;
     const regex = /\[(\d+)\]/g;
     let processedMessage = message.content;
 
@@ -119,9 +118,7 @@ const MessageTabs = ({
       message.sources.length > 0
     ) {
       setParsedMessage(
-        processedMessage.replace(
+        processedMessage.replace(regex, (_, capturedContent: string) => {
-          citationRegex,
-          (_, capturedContent: string) => {
           const numbers = capturedContent
             .split(',')
             .map((numStr) => numStr.trim());
@@ -146,8 +143,7 @@ const MessageTabs = ({
             .join('');
 
           return linksHtml;
-          },
+        }),
-        ),
       );
       setSpeechMessage(message.content.replace(regex, ''));
       return;

src/lib/search/simplifiedAgent.ts

@@ -58,11 +58,8 @@ function normalizeUsageMetadata(usageData: any): {
 }
 
 /**
- * Simplified Agent using createReactAgent
+ * SimplifiedAgent class that provides a streamlined interface for creating and managing an AI agent
- *
+ * with customizable focus modes and tools.
- * This agent replaces the complex LangGraph supervisor pattern with a single
- * tool-calling agent that handles analysis and synthesis internally while
- * using specialized tools for search, file processing, and URL summarization.
  */
 export class SimplifiedAgent {
   private llm: BaseChatModel;
@@ -95,7 +92,6 @@ export class SimplifiedAgent {
     // Select appropriate tools based on focus mode and available files
     const tools = this.getToolsForFocusMode(focusMode, fileIds);
 
-    // Create the enhanced system prompt that includes analysis and synthesis instructions
     const enhancedSystemPrompt = this.createEnhancedSystemPrompt(
       focusMode,
       fileIds,
@@ -159,9 +155,6 @@ export class SimplifiedAgent {
     }
   }
 
-  /**
-   * Create enhanced system prompt that includes analysis and synthesis capabilities
-   */
   private createEnhancedSystemPrompt(
     focusMode: string,
     fileIds: string[] = [],
@@ -348,15 +341,16 @@ Your task is to provide answers that are:
     - Passing true is **required** to include images or links within the page content.
   - You will receive a summary of the content from each URL if the content of the page is long. If the content of the page is short, you will receive the full content.
   - You may request up to 5 URLs per turn.
+  - If you recieve a request to summarize a specific URL you **must** use this tool to retrieve it.
 5. **Analyze**: Examine the retrieved information for relevance, accuracy, and completeness.
-  - If you have sufficient information, you can move on to the synthesis stage.
+  - If you have sufficient information, you can move on to the respond stage.
   - If you need to gather more information, consider revisiting the search or supplement stages.${
     fileIds.length > 0
       ? `
   - Consider both web search results and file content when analyzing information completeness.`
       : ''
   }
-6. **Synthesize**: Combine all information into a coherent, well-cited response
+6. **Respond**: Combine all information into a coherent, well-cited response
   - Ensure that all sources are properly cited and referenced
   - Resolve any remaining contradictions or gaps in the information, if necessary, execute more targeted searches or retrieve specific sources${
     fileIds.length > 0
@@ -457,13 +451,14 @@ Your task is to provide answers that are:
   - You will receive relevant excerpts from documents that match your search criteria.
   - Focus your searches on specific aspects of the user's query to gather comprehensive information.
 3. **Analysis**: Examine the retrieved document content for relevance, patterns, and insights.
-  - If you have sufficient information from the documents, you can move on to the synthesis stage.
+  - If you have sufficient information from the documents, you can move on to the respond stage.
   - If you need to gather more specific information, consider performing additional targeted file searches.
   - Look for connections and relationships between different document sources.
-4. **Synthesize**: Combine all document insights into a coherent, well-cited response
+4. **Respond**: Combine all document insights into a coherent, well-cited response
   - Ensure that all sources are properly cited and referenced
   - Resolve any contradictions or gaps in the document information
   - Provide comprehensive analysis based on the available document content
+  - Only respond with your final answer once you've gathered all relevant information and are done with tool use
 
 ## Current Context
 - Today's Date: ${formatDateForLLM(new Date())}