Merge pull request #63 from MODSetter/dev

feat(ui): performance fixes

surfsense_web/app/dashboard/[search_space_id]/researcher/[chat_id]/page.tsx

@@ -240,7 +240,7 @@ const SourcesDialogContent = ({
 const ChatPage = () => {
   const [token, setToken] = React.useState<string | null>(null);
   const [activeTab, setActiveTab] = useState("");
-  const [dialogOpen, setDialogOpen] = useState(false);
+  const [dialogOpenId, setDialogOpenId] = useState<number | null>(null);
   const [sourcesPage, setSourcesPage] = useState(1);
   const [expandedSources, setExpandedSources] = useState(false);
   const [canScrollLeft, setCanScrollLeft] = useState(false);
@@ -260,6 +260,13 @@ const ChatPage = () => {
 
   const { search_space_id, chat_id } = useParams();
 
+  // Function to scroll terminal to bottom
+  const scrollTerminalToBottom = () => {
+    if (terminalMessagesRef.current) {
+      terminalMessagesRef.current.scrollTop = terminalMessagesRef.current.scrollHeight;
+    }
+  };
+
   // Get token from localStorage on client side only
   React.useEffect(() => {
     setToken(localStorage.getItem('surfsense_bearer_token'));
@@ -469,54 +476,60 @@ const ChatPage = () => {
     updateChat();
   }, [messages, status, chat_id, researchMode, selectedConnectors, search_space_id]);
 
-  // Log messages whenever they update and extract annotations from the latest assistant message if available
+  // Memoize connector sources to prevent excessive re-renders
-  useEffect(() => {
+  const processedConnectorSources = React.useMemo(() => {
-    console.log('Messages updated:', messages);
+    if (messages.length === 0) return connectorSources;
     
-    // Extract annotations from the latest assistant message if available
+    // Only process when we have a complete message (not streaming)
+    if (status !== 'ready') return connectorSources;
+    
+    // Find the latest assistant message
     const assistantMessages = messages.filter(msg => msg.role === 'assistant');
-    if (assistantMessages.length > 0) {
+    if (assistantMessages.length === 0) return connectorSources;
-      const latestAssistantMessage = assistantMessages[assistantMessages.length - 1];
-      if (latestAssistantMessage?.annotations) {
-        const annotations = latestAssistantMessage.annotations as any[];
     
-        // Debug log to track streaming annotations
+    const latestAssistantMessage = assistantMessages[assistantMessages.length - 1];
-        if (process.env.NODE_ENV === 'development') {
+    if (!latestAssistantMessage?.annotations) return connectorSources;
-          console.log('Streaming annotations:', annotations);
     
-          // Log counts of each annotation type
+    // Find the latest SOURCES annotation
-          const terminalInfoCount = annotations.filter(a => a.type === 'TERMINAL_INFO').length;
+    const annotations = latestAssistantMessage.annotations as any[];
-          const sourcesCount = annotations.filter(a => a.type === 'SOURCES').length;
+    const sourcesAnnotations = annotations.filter(a => a.type === 'SOURCES');
-          const answerCount = annotations.filter(a => a.type === 'ANSWER').length;
     
-          console.log(`Annotation counts - Terminal: ${terminalInfoCount}, Sources: ${sourcesCount}, Answer: ${answerCount}`);
+    if (sourcesAnnotations.length === 0) return connectorSources;
-        }
     
-        // Process SOURCES annotation - get the last one to ensure we have the latest
+    const latestSourcesAnnotation = sourcesAnnotations[sourcesAnnotations.length - 1];
-        const sourcesAnnotations = annotations.filter(
+    if (!latestSourcesAnnotation.content) return connectorSources;
-          (annotation) => annotation.type === 'SOURCES'
-        );
     
-        if (sourcesAnnotations.length > 0) {
+    // Use this content if it differs from current
-          // Get the last SOURCES annotation to ensure we have the most recent one
+    return latestSourcesAnnotation.content;
-          const latestSourcesAnnotation = sourcesAnnotations[sourcesAnnotations.length - 1];
+  }, [messages, status, connectorSources]);
-          if (latestSourcesAnnotation.content) {
-            setConnectorSources(latestSourcesAnnotation.content);
-          }
-        }
   
-        // Check for terminal info annotations and scroll terminal to bottom if they exist
+  // Update connector sources when processed value changes
-        const terminalInfoAnnotations = annotations.filter(
+  useEffect(() => {
-          (annotation) => annotation.type === 'TERMINAL_INFO'
+    if (processedConnectorSources !== connectorSources) {
-        );
+      setConnectorSources(processedConnectorSources);
-
-        if (terminalInfoAnnotations.length > 0) {
-          // Schedule scrolling after the DOM has been updated
-          setTimeout(scrollTerminalToBottom, 100);
-        }
-      }
     }
-  }, [messages]);
+  }, [processedConnectorSources, connectorSources]);
+  
+  // Check and scroll terminal when terminal info is available
+  useEffect(() => {
+    if (messages.length === 0 || status !== 'ready') return;
+    
+    // Find the latest assistant message
+    const assistantMessages = messages.filter(msg => msg.role === 'assistant');
+    if (assistantMessages.length === 0) return;
+    
+    const latestAssistantMessage = assistantMessages[assistantMessages.length - 1];
+    if (!latestAssistantMessage?.annotations) return;
+    
+    // Check for terminal info annotations
+    const annotations = latestAssistantMessage.annotations as any[];
+    const terminalInfoAnnotations = annotations.filter(a => a.type === 'TERMINAL_INFO');
+    
+    if (terminalInfoAnnotations.length > 0) {
+      // Schedule scrolling after the DOM has been updated
+      setTimeout(scrollTerminalToBottom, 100);
+    }
+  }, [messages, status]);
 
   // Custom handleSubmit function to include selected connectors and answer type
   const handleSubmit = (e: React.FormEvent) => {
@@ -543,25 +556,23 @@ const ChatPage = () => {
     messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
   };
 
-  // Function to scroll terminal to bottom
-  const scrollTerminalToBottom = () => {
-    if (terminalMessagesRef.current) {
-      terminalMessagesRef.current.scrollTop = terminalMessagesRef.current.scrollHeight;
-    }
-  };
-
   // Scroll to bottom when messages change
   useEffect(() => {
     scrollToBottom();
   }, [messages]);
 
-  // Set activeTab when connectorSources change
+  // Set activeTab when connectorSources change using a memoized value
-  useEffect(() => {
+  const activeTabValue = React.useMemo(() => {
-    if (connectorSources.length > 0) {
+    return connectorSources.length > 0 ? connectorSources[0].type : "";
-      setActiveTab(connectorSources[0].type);
-    }
   }, [connectorSources]);
   
+  // Update activeTab when the memoized value changes
+  useEffect(() => {
+    if (activeTabValue && activeTabValue !== activeTab) {
+      setActiveTab(activeTabValue);
+    }
+  }, [activeTabValue, activeTab]);
+
   // Scroll terminal to bottom when expanded
   useEffect(() => {
     if (terminalExpanded) {
@@ -617,7 +628,7 @@ const ChatPage = () => {
   };
 
   // Function to get a citation source by ID
-  const getCitationSource = (citationId: number, messageIndex?: number): Source | null => {
+  const getCitationSource = React.useCallback((citationId: number, messageIndex?: number): Source | null => {
     if (!messages || messages.length === 0) return null;
 
     // If no specific message index is provided, use the latest assistant message
@@ -699,7 +710,7 @@ const ChatPage = () => {
 
       return foundSource || null;
     }
-  };
+  }, [messages]);
 
   return (
     <>
@@ -900,7 +911,7 @@ const ChatPage = () => {
                               ))}
 
                               {connector.sources.length > INITIAL_SOURCES_DISPLAY && (
-                                <Dialog open={dialogOpen && activeTab === connector.type} onOpenChange={(open) => setDialogOpen(open)}>
+                                <Dialog open={dialogOpenId === connector.id} onOpenChange={(open) => setDialogOpenId(open ? connector.id : null)}>
                                   <DialogTrigger asChild>
                                     <Button variant="ghost" className="w-full text-sm text-gray-500 dark:text-gray-400">
                                       Show {connector.sources.length - INITIAL_SOURCES_DISPLAY} More Sources

surfsense_web/components/chat/Citation.tsx

@@ -20,7 +20,7 @@ type CitationProps = {
 /**
  * Citation component to handle individual citations
  */
-export const Citation = ({ citationId, citationText, position, source }: CitationProps) => {
+export const Citation = React.memo(({ citationId, citationText, position, source }: CitationProps) => {
   const [open, setOpen] = useState(false);
   const citationKey = `citation-${citationId}-${position}`;
   
@@ -38,37 +38,41 @@ export const Citation = ({ citationId, citationText, position, source }: Citatio
             </span>
           </sup>
         </DropdownMenuTrigger>
-        <DropdownMenuContent align="start" className="w-80 p-0">
+        {open && (
-          <Card className="border-0 shadow-none">
+          <DropdownMenuContent align="start" className="w-80 p-0" forceMount>
-            <div className="p-3 flex items-start gap-3">
+            <Card className="border-0 shadow-none">
-              <div className="flex-shrink-0 w-7 h-7 flex items-center justify-center bg-muted rounded-full">
+              <div className="p-3 flex items-start gap-3">
-                {getConnectorIcon(source.connectorType || '')}
+                <div className="flex-shrink-0 w-7 h-7 flex items-center justify-center bg-muted rounded-full">
-              </div>
+                  {getConnectorIcon(source.connectorType || '')}
-              <div className="flex-1">
-                <div className="flex items-center gap-2 mb-1">
-                  <h3 className="font-medium text-sm text-card-foreground">{source.title}</h3>
                 </div>
-                <p className="text-sm text-muted-foreground mt-0.5">{source.description}</p>
+                <div className="flex-1">
-                <div className="mt-2 flex items-center text-xs text-muted-foreground">
+                  <div className="flex items-center gap-2 mb-1">
-                  <span className="truncate max-w-[200px]">{source.url}</span>
+                    <h3 className="font-medium text-sm text-card-foreground">{source.title}</h3>
+                  </div>
+                  <p className="text-sm text-muted-foreground mt-0.5">{source.description}</p>
+                  <div className="mt-2 flex items-center text-xs text-muted-foreground">
+                    <span className="truncate max-w-[200px]">{source.url}</span>
+                  </div>
                 </div>
+                <Button 
+                  variant="ghost" 
+                  size="icon" 
+                  className="h-7 w-7 rounded-full"
+                  onClick={() => window.open(source.url, '_blank', 'noopener,noreferrer')}
+                  title="Open in new tab"
+                >
+                  <ExternalLink className="h-3.5 w-3.5" />
+                </Button>
               </div>
-              <Button 
+            </Card>
-                variant="ghost" 
+          </DropdownMenuContent>
-                size="icon" 
+        )}
-                className="h-7 w-7 rounded-full"
-                onClick={() => window.open(source.url, '_blank')}
-                title="Open in new tab"
-              >
-                <ExternalLink className="h-3.5 w-3.5" />
-              </Button>
-            </div>
-          </Card>
-        </DropdownMenuContent>
       </DropdownMenu>
     </span>
   );
-};
+});
+
+Citation.displayName = 'Citation';
 
 /**
  * Function to render text with citations

surfsense_web/components/markdown-viewer.tsx

@@ -1,4 +1,4 @@
-import React from "react";
+import React, { useMemo } from "react";
 import ReactMarkdown from "react-markdown";
 import rehypeRaw from "rehype-raw";
 import rehypeSanitize from "rehype-sanitize";
@@ -14,82 +14,87 @@ interface MarkdownViewerProps {
 }
 
 export function MarkdownViewer({ content, className, getCitationSource }: MarkdownViewerProps) {
+  // Memoize the markdown components to prevent unnecessary re-renders
+  const components = useMemo(() => {
+    return {
+      // Define custom components for markdown elements
+      p: ({node, children, ...props}: any) => {
+        // If there's no getCitationSource function, just render normally
+        if (!getCitationSource) {
+          return <p className="my-2" {...props}>{children}</p>;
+        }
+        
+        // Process citations within paragraph content
+        return <p className="my-2" {...props}>{processCitationsInReactChildren(children, getCitationSource)}</p>;
+      },
+      a: ({node, children, ...props}: any) => {
+        // Process citations within link content if needed
+        const processedChildren = getCitationSource 
+          ? processCitationsInReactChildren(children, getCitationSource) 
+          : children;
+        return <a className="text-primary hover:underline" {...props}>{processedChildren}</a>;
+      },
+      li: ({node, children, ...props}: any) => {
+        // Process citations within list item content
+        const processedChildren = getCitationSource 
+          ? processCitationsInReactChildren(children, getCitationSource) 
+          : children;
+        return <li {...props}>{processedChildren}</li>;
+      },
+      ul: ({node, ...props}: any) => <ul className="list-disc pl-5 my-2" {...props} />,
+      ol: ({node, ...props}: any) => <ol className="list-decimal pl-5 my-2" {...props} />,
+      h1: ({node, children, ...props}: any) => {
+        const processedChildren = getCitationSource 
+          ? processCitationsInReactChildren(children, getCitationSource) 
+          : children;
+        return <h1 className="text-2xl font-bold mt-6 mb-2" {...props}>{processedChildren}</h1>;
+      },
+      h2: ({node, children, ...props}: any) => {
+        const processedChildren = getCitationSource 
+          ? processCitationsInReactChildren(children, getCitationSource) 
+          : children;
+        return <h2 className="text-xl font-bold mt-5 mb-2" {...props}>{processedChildren}</h2>;
+      },
+      h3: ({node, children, ...props}: any) => {
+        const processedChildren = getCitationSource 
+          ? processCitationsInReactChildren(children, getCitationSource) 
+          : children;
+        return <h3 className="text-lg font-bold mt-4 mb-2" {...props}>{processedChildren}</h3>;
+      },
+      h4: ({node, children, ...props}: any) => {
+        const processedChildren = getCitationSource 
+          ? processCitationsInReactChildren(children, getCitationSource) 
+          : children;
+        return <h4 className="text-base font-bold mt-3 mb-1" {...props}>{processedChildren}</h4>;
+      },
+      blockquote: ({node, ...props}: any) => <blockquote className="border-l-4 border-muted pl-4 italic my-2" {...props} />,
+      hr: ({node, ...props}: any) => <hr className="my-4 border-muted" {...props} />,
+      img: ({node, ...props}: any) => <img className="max-w-full h-auto my-4 rounded" {...props} />,
+      table: ({node, ...props}: any) => <div className="overflow-x-auto my-4"><table className="min-w-full divide-y divide-border" {...props} /></div>,
+      th: ({node, ...props}: any) => <th className="px-3 py-2 text-left font-medium bg-muted" {...props} />,
+      td: ({node, ...props}: any) => <td className="px-3 py-2 border-t border-border" {...props} />,
+      code: ({node, className, children, ...props}: any) => {
+        const match = /language-(\w+)/.exec(className || '');
+        const isInline = !match;
+        return isInline 
+          ? <code className="bg-muted px-1 py-0.5 rounded text-xs" {...props}>{children}</code>
+          : (
+            <div className="relative my-4">
+              <pre className="bg-muted p-4 rounded-md overflow-x-auto">
+                <code className="text-xs" {...props}>{children}</code>
+              </pre>
+            </div>
+          );
+      }
+    };
+  }, [getCitationSource]);
+
   return (
     <div className={cn("prose prose-sm dark:prose-invert max-w-none", className)}>
       <ReactMarkdown
         rehypePlugins={[rehypeRaw, rehypeSanitize]}
         remarkPlugins={[remarkGfm]}
-        components={{
+        components={components}
-          // Define custom components for markdown elements
-          p: ({node, children, ...props}) => {
-            // If there's no getCitationSource function, just render normally
-            if (!getCitationSource) {
-              return <p className="my-2" {...props}>{children}</p>;
-            }
-            
-            // Process citations within paragraph content
-            return <p className="my-2" {...props}>{processCitationsInReactChildren(children, getCitationSource)}</p>;
-          },
-          a: ({node, children, ...props}) => {
-            // Process citations within link content if needed
-            const processedChildren = getCitationSource 
-              ? processCitationsInReactChildren(children, getCitationSource) 
-              : children;
-            return <a className="text-primary hover:underline" {...props}>{processedChildren}</a>;
-          },
-          li: ({node, children, ...props}) => {
-            // Process citations within list item content
-            const processedChildren = getCitationSource 
-              ? processCitationsInReactChildren(children, getCitationSource) 
-              : children;
-            return <li {...props}>{processedChildren}</li>;
-          },
-          ul: ({node, ...props}) => <ul className="list-disc pl-5 my-2" {...props} />,
-          ol: ({node, ...props}) => <ol className="list-decimal pl-5 my-2" {...props} />,
-          h1: ({node, children, ...props}) => {
-            const processedChildren = getCitationSource 
-              ? processCitationsInReactChildren(children, getCitationSource) 
-              : children;
-            return <h1 className="text-2xl font-bold mt-6 mb-2" {...props}>{processedChildren}</h1>;
-          },
-          h2: ({node, children, ...props}) => {
-            const processedChildren = getCitationSource 
-              ? processCitationsInReactChildren(children, getCitationSource) 
-              : children;
-            return <h2 className="text-xl font-bold mt-5 mb-2" {...props}>{processedChildren}</h2>;
-          },
-          h3: ({node, children, ...props}) => {
-            const processedChildren = getCitationSource 
-              ? processCitationsInReactChildren(children, getCitationSource) 
-              : children;
-            return <h3 className="text-lg font-bold mt-4 mb-2" {...props}>{processedChildren}</h3>;
-          },
-          h4: ({node, children, ...props}) => {
-            const processedChildren = getCitationSource 
-              ? processCitationsInReactChildren(children, getCitationSource) 
-              : children;
-            return <h4 className="text-base font-bold mt-3 mb-1" {...props}>{processedChildren}</h4>;
-          },
-          blockquote: ({node, ...props}) => <blockquote className="border-l-4 border-muted pl-4 italic my-2" {...props} />,
-          hr: ({node, ...props}) => <hr className="my-4 border-muted" {...props} />,
-          img: ({node, ...props}) => <img className="max-w-full h-auto my-4 rounded" {...props} />,
-          table: ({node, ...props}) => <div className="overflow-x-auto my-4"><table className="min-w-full divide-y divide-border" {...props} /></div>,
-          th: ({node, ...props}) => <th className="px-3 py-2 text-left font-medium bg-muted" {...props} />,
-          td: ({node, ...props}) => <td className="px-3 py-2 border-t border-border" {...props} />,
-          code: ({node, className, children, ...props}: any) => {
-            const match = /language-(\w+)/.exec(className || '');
-            const isInline = !match;
-            return isInline 
-              ? <code className="bg-muted px-1 py-0.5 rounded text-xs" {...props}>{children}</code>
-              : (
-                <div className="relative my-4">
-                  <pre className="bg-muted p-4 rounded-md overflow-x-auto">
-                    <code className="text-xs" {...props}>{children}</code>
-                  </pre>
-                </div>
-              );
-          }
-        }}
       >
         {content}
       </ReactMarkdown>
@@ -98,7 +103,7 @@ export function MarkdownViewer({ content, className, getCitationSource }: Markdo
 }
 
 // Helper function to process citations within React children
-function processCitationsInReactChildren(children: React.ReactNode, getCitationSource: (id: number) => Source | null): React.ReactNode {
+const processCitationsInReactChildren = (children: React.ReactNode, getCitationSource: (id: number) => Source | null): React.ReactNode => {
   // If children is not an array or string, just return it
   if (!children || (typeof children !== 'string' && !Array.isArray(children))) {
     return children;
@@ -120,10 +125,10 @@ function processCitationsInReactChildren(children: React.ReactNode, getCitationS
   }
   
   return children;
-}
+};
 
 // Process citation references in text content
-function processCitationsInText(text: string, getCitationSource: (id: number) => Source | null): React.ReactNode[] {
+const processCitationsInText = (text: string, getCitationSource: (id: number) => Source | null): React.ReactNode[] => {
   // Use improved regex to catch citation numbers more reliably
   // This will match patterns like [1], [42], etc. including when they appear at the end of a line or sentence
   const citationRegex = /\[(\d+)\]/g;
@@ -132,13 +137,7 @@ function processCitationsInText(text: string, getCitationSource: (id: number) =>
   let match;
   let position = 0;
   
-  // Debug log for troubleshooting
-  console.log("Processing citations in text:", text);
-  
   while ((match = citationRegex.exec(text)) !== null) {
-    // Log each match for debugging
-    console.log("Citation match found:", match[0], "at index", match.index);
-    
     // Add text before the citation
     if (match.index > lastIndex) {
       parts.push(text.substring(lastIndex, match.index));
@@ -148,9 +147,6 @@ function processCitationsInText(text: string, getCitationSource: (id: number) =>
     const citationId = parseInt(match[1], 10);
     const source = getCitationSource(citationId);
     
-    // Log the citation details
-    console.log("Citation ID:", citationId, "Source:", source ? "found" : "not found");
-    
     parts.push(
       <Citation 
         key={`citation-${citationId}-${position}`}
@@ -171,4 +167,4 @@ function processCitationsInText(text: string, getCitationSource: (id: number) =>
   }
   
   return parts;
-} 
+};