AI Assistant Performance (#32792)

apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx

@@ -141,16 +141,10 @@ export const AIAssistant = ({
   const { mutate: updateOrganization, isLoading: isUpdating } = useOrganizationUpdateMutation()
 
   const messages = useMemo(() => {
-    const merged = [
+    return [
       ...chatMessages,
       ...(assistantError !== undefined && lastSentMessage !== undefined ? [lastSentMessage] : []),
     ]
-
-    return merged.sort(
-      (a, b) =>
-        (a.createdAt?.getTime() ?? 0) - (b.createdAt?.getTime() ?? 0) ||
-        a.role.localeCompare(b.role)
-    )
   }, [chatMessages, assistantError, lastSentMessage])
 
   const renderedMessages = useMemo(
@@ -160,7 +154,7 @@ export const AIAssistant = ({
           <MemoizedMessage
             key={message.id}
             message={message}
-            isLoading={isChatLoading && message === messages[messages.length - 1]}
+            isLoading={isChatLoading && message.id === messages[messages.length - 1].id}
           />
         )
       }),
@@ -313,12 +307,12 @@ export const AIAssistant = ({
             </div>
           )}
           {hasMessages ? (
-            <motion.div className="w-full p-5">
+            <div className="w-full p-5">
               {renderedMessages}
               {(last(messages)?.role === 'user' || last(messages)?.content?.length === 0) && (
                 <div className="flex gap-4 w-auto overflow-hidden">
                   <AiIconAnimation size={20} className="text-foreground-muted shrink-0" />
-                  <motion.div className="text-foreground-lighter text-sm flex gap-1.5 items-center">
+                  <div className="text-foreground-lighter text-sm flex gap-1.5 items-center">
                     <span>Thinking</span>
                     <div className="flex gap-1">
                       <motion.span
@@ -340,11 +334,11 @@ export const AIAssistant = ({
                         .
                       </motion.span>
                     </div>
-                  </motion.div>
+                  </div>
                 </div>
               )}
               <div className="h-1" />
-            </motion.div>
+            </div>
           ) : suggestions ? (
             <div className="w-full h-full px-8 py-0 flex flex-col flex-1 justify-end">
               <h3 className="text-foreground-light font-mono text-sm uppercase mb-3">

apps/studio/components/ui/AIAssistantPanel/Message.tsx

@@ -1,12 +1,28 @@
-import { motion } from 'framer-motion'
 import { User } from 'lucide-react'
-import { PropsWithChildren } from 'react'
+import {
+  PropsWithChildren,
+  memo,
+  useMemo,
+  useEffect,
+  createContext,
+  useContext,
+  ReactNode,
+} from 'react'
 import ReactMarkdown from 'react-markdown'
 import remarkGfm from 'remark-gfm'
+import { Components } from 'react-markdown/lib/ast-to-react'
 
-import { AiIconAnimation, cn, CodeBlock, markdownComponents, WarningIcon } from 'ui'
+import {
+  AiIconAnimation,
+  cn,
+  CodeBlock,
+  markdownComponents,
+  WarningIcon,
+  type CodeBlockLang,
+} from 'ui'
 import { QueryBlock } from '../QueryBlock/QueryBlock'
 import CollapsibleCodeBlock from './CollapsibleCodeBlock'
+import { DebouncedComponent } from '../DebouncedComponent'
 
 interface MessageProps {
   role: 'function' | 'system' | 'user' | 'assistant' | 'data' | 'tool'
@@ -25,6 +41,150 @@ type AssistantSnippetProps = {
   yAxis?: string
 }
 
+interface MessageContextType {
+  isLoading: boolean
+  readOnly?: boolean
+}
+
+const MessageContext = createContext<MessageContextType>({ isLoading: false })
+
+const MemoizedQueryBlock = memo(
+  ({
+    sql,
+    title,
+    xAxis,
+    yAxis,
+    isChart,
+    isLoading,
+    runQuery,
+  }: {
+    sql: string
+    title: string
+    xAxis?: string
+    yAxis?: string
+    isChart: boolean
+    isLoading: boolean
+    runQuery: boolean
+  }) => (
+    <DebouncedComponent
+      delay={500}
+      value={sql}
+      fallback={
+        <div className="bg-surface-100 border-overlay rounded border shadow-sm px-3 py-2 text-xs">
+          Writing SQL...
+        </div>
+      }
+    >
+      <QueryBlock
+        lockColumns
+        label={title}
+        sql={sql}
+        chartConfig={{
+          type: 'bar',
+          cumulative: false,
+          xKey: xAxis ?? '',
+          yKey: yAxis ?? '',
+        }}
+        isChart={isChart}
+        isLoading={isLoading}
+        runQuery={runQuery}
+      />
+    </DebouncedComponent>
+  )
+)
+MemoizedQueryBlock.displayName = 'MemoizedQueryBlock'
+
+const MarkdownPre = ({ children }: any) => {
+  const { isLoading, readOnly } = useContext(MessageContext)
+
+  const language = children[0].props.className?.replace('language-', '') || 'sql'
+  const rawSql = language === 'sql' ? children[0].props.children : undefined
+  const formatted = (rawSql || [''])[0]
+  const propsMatch = formatted.match(/--\s*props:\s*(\{[^}]+\})/)
+
+  const snippetProps: AssistantSnippetProps = propsMatch ? JSON.parse(propsMatch[1]) : {}
+  const { xAxis, yAxis } = snippetProps
+  const title = snippetProps.title || 'SQL Query'
+  const isChart = snippetProps.isChart === 'true'
+  const runQuery = snippetProps.runQuery === 'true'
+  const sql = formatted?.replace(/--\s*props:\s*\{[^}]+\}/, '').trim()
+
+  return (
+    <div className="w-auto -ml-[36px] overflow-x-hidden">
+      {language === 'sql' ? (
+        readOnly ? (
+          <CollapsibleCodeBlock
+            value={children[0].props.children[0]}
+            language="sql"
+            hideLineNumbers
+          />
+        ) : (
+          <MemoizedQueryBlock
+            sql={sql}
+            title={title}
+            xAxis={xAxis}
+            yAxis={yAxis}
+            isChart={isChart}
+            isLoading={isLoading}
+            runQuery={runQuery}
+          />
+        )
+      ) : (
+        <CodeBlock
+          hideLineNumbers
+          value={children[0].props.children[0]}
+          language={language as CodeBlockLang}
+          className={cn(
+            'max-h-96 max-w-none block border rounded !bg-transparent !py-3 !px-3.5 prose dark:prose-dark text-foreground',
+            '[&>code]:m-0 [&>code>span]:flex [&>code>span]:flex-wrap [&>code]:block [&>code>span]:text-foreground'
+          )}
+        />
+      )}
+    </div>
+  )
+}
+
+const OrderedList = memo(({ children }: { children: ReactNode }) => (
+  <ol className="flex flex-col gap-y-4">{children}</ol>
+))
+OrderedList.displayName = 'OrderedList'
+
+const ListItem = memo(({ children }: { children: ReactNode }) => (
+  <li className="[&>pre]:mt-2">{children}</li>
+))
+ListItem.displayName = 'ListItem'
+
+const Heading3 = memo(({ children }: { children: ReactNode }) => (
+  <h3 className="underline">{children}</h3>
+))
+Heading3.displayName = 'Heading3'
+
+const InlineCode = memo(({ className, children }: { className?: string; children: ReactNode }) => (
+  <code className={cn('text-xs', className)}>{children}</code>
+))
+InlineCode.displayName = 'InlineCode'
+
+const Link = memo(({ href, children }: { href?: string; children: ReactNode }) => (
+  <a
+    target="_blank"
+    rel="noopener noreferrer"
+    href={href}
+    className="underline transition underline-offset-2 decoration-foreground-lighter hover:decoration-foreground text-foreground"
+  >
+    {children}
+  </a>
+))
+Link.displayName = 'Link'
+
+const baseMarkdownComponents: Partial<Components> = {
+  ol: OrderedList,
+  li: ListItem,
+  h3: Heading3,
+  code: InlineCode,
+  a: Link,
+  pre: MarkdownPre,
+}
+
 export const Message = function Message({
   role,
   content,
@@ -34,124 +194,56 @@ export const Message = function Message({
   action = null,
   variant = 'default',
 }: PropsWithChildren<MessageProps>) {
+  useEffect(() => {
+    return () => {
+      console.log('unmounting parent')
+    }
+  }, [])
+
   const isUser = role === 'user'
 
+  const allMarkdownComponents = useMemo(
+    () => ({
+      ...markdownComponents,
+      ...baseMarkdownComponents,
+    }),
+    []
+  )
+
   if (!content) return null
 
   return (
-    <motion.div
-      layout="position"
-      initial={{ opacity: 0 }}
-      animate={{ opacity: 1 }}
-      className={cn(
-        'mb-5 text-foreground-light text-sm',
-        isUser && 'text-foreground',
-        variant === 'warning' && 'bg-warning-200'
-      )}
-    >
-      {children}
-
-      {variant === 'warning' && <WarningIcon className="w-6 h-6" />}
-
-      {action}
-
-      <div className="flex gap-4 w-auto overflow-hidden">
-        {isUser ? (
-          <figure className="w-5 h-5 shrink-0 bg-foreground rounded-full flex items-center justify-center">
-            <User size={16} strokeWidth={1.5} className="text-background" />
-          </figure>
-        ) : (
-          <AiIconAnimation size={20} className="text-foreground-muted shrink-0" />
+    <MessageContext.Provider value={{ isLoading, readOnly }}>
+      <div
+        className={cn(
+          'mb-5 text-foreground-light text-sm',
+          isUser && 'text-foreground',
+          variant === 'warning' && 'bg-warning-200'
         )}
-        <ReactMarkdown
-          className="space-y-5 flex-1 [&>*>code]:text-xs [&>*>*>code]:text-xs min-w-0 [&_li]:space-y-4"
-          remarkPlugins={[remarkGfm]}
-          components={{
-            ...markdownComponents,
-            pre: (props: any) => {
-              const language = props.children[0].props.className?.replace('language-', '') || 'sql'
+      >
+        {children}
 
-              const rawSql = language === 'sql' ? props.children[0].props.children : undefined
-              const formatted = (rawSql || [''])[0]
-              const propsMatch = formatted.match(/--\s*props:\s*(\{[^}]+\})/)
+        {variant === 'warning' && <WarningIcon className="w-6 h-6" />}
 
-              const snippetProps: AssistantSnippetProps = propsMatch
-                ? JSON.parse(propsMatch[1])
-                : {}
-              const { xAxis, yAxis } = snippetProps
-              const title = snippetProps.title || 'SQL Query'
-              const isChart = snippetProps.isChart === 'true'
-              const runQuery = snippetProps.runQuery === 'true'
-              const sql = formatted?.replace(/--\s*props:\s*\{[^}]+\}/, '').trim()
+        {action}
 
-              return (
-                <div className="w-auto -ml-[36px] overflow-x-hidden">
-                  {language === 'sql' ? (
-                    readOnly ? (
-                      <CollapsibleCodeBlock
-                        value={props.children[0].props.children[0]}
-                        language="sql"
-                        hideLineNumbers
-                      />
-                    ) : (
-                      <QueryBlock
-                        lockColumns
-                        label={title}
-                        sql={sql}
-                        chartConfig={{
-                          type: 'bar',
-                          cumulative: false,
-                          xKey: xAxis ?? '',
-                          yKey: yAxis ?? '',
-                        }}
-                        isChart={isChart}
-                        isLoading={isLoading}
-                        runQuery={runQuery}
-                      />
-                    )
-                  ) : (
-                    <CodeBlock
-                      hideLineNumbers
-                      value={props.children[0].props.children}
-                      language={language}
-                      className={cn(
-                        'max-h-96 max-w-none block border rounded !bg-transparent !py-3 !px-3.5 prose dark:prose-dark text-foreground',
-                        '[&>code]:m-0 [&>code>span]:flex [&>code>span]:flex-wrap [&>code]:block [&>code>span]:text-foreground'
-                      )}
-                    />
-                  )}
-                </div>
-              )
-            },
-            ol: (props: any) => {
-              return <ol className="flex flex-col gap-y-4">{props.children}</ol>
-            },
-            li: (props: any) => {
-              return <li className="[&>pre]:mt-2">{props.children}</li>
-            },
-            h3: (props: any) => {
-              return <h3 className="underline">{props.children}</h3>
-            },
-            code: (props: any) => {
-              return <code className={cn('text-xs', props.className)}>{props.children}</code>
-            },
-            a: (props: any) => {
-              return (
-                <a
-                  target="_blank"
-                  rel="noopener noreferrer"
-                  href={props.href}
-                  className="underline transition underline-offset-2 decoration-foreground-lighter hover:decoration-foreground text-foreground"
-                >
-                  {props.children}
-                </a>
-              )
-            },
-          }}
-        >
-          {content}
-        </ReactMarkdown>
+        <div className="flex gap-4 w-auto overflow-hidden">
+          {isUser ? (
+            <figure className="w-5 h-5 shrink-0 bg-foreground rounded-full flex items-center justify-center">
+              <User size={16} strokeWidth={1.5} className="text-background" />
+            </figure>
+          ) : (
+            <AiIconAnimation size={20} className="text-foreground-muted shrink-0" />
+          )}
+          <ReactMarkdown
+            className="space-y-5 flex-1 [&>*>code]:text-xs [&>*>*>code]:text-xs min-w-0 [&_li]:space-y-4"
+            remarkPlugins={[remarkGfm]}
+            components={allMarkdownComponents}
+          >
+            {content}
+          </ReactMarkdown>
+        </div>
       </div>
-    </motion.div>
+    </MessageContext.Provider>
   )
 }

apps/studio/components/ui/DebouncedComponent.tsx

@@ -0,0 +1,44 @@
+import { useEffect, useRef, useState } from 'react'
+
+interface DebouncedComponentProps {
+  value: any
+  delay?: number
+  fallback?: React.ReactNode
+  children: React.ReactNode
+}
+
+export function DebouncedComponent({
+  value,
+  delay = 500,
+  fallback = <div className="text-sm">Loading...</div>,
+  children,
+}: DebouncedComponentProps) {
+  const [shouldRender, setShouldRender] = useState(false)
+  const timeoutRef = useRef<NodeJS.Timeout>()
+  const prevValueRef = useRef(value)
+  const isInitialMount = useRef(true)
+
+  useEffect(() => {
+    if (isInitialMount.current || prevValueRef.current !== value) {
+      setShouldRender(false)
+      prevValueRef.current = value
+
+      if (timeoutRef.current) {
+        clearTimeout(timeoutRef.current)
+      }
+
+      timeoutRef.current = setTimeout(() => {
+        setShouldRender(true)
+        isInitialMount.current = false
+      }, delay)
+    }
+
+    return () => {
+      if (timeoutRef.current) {
+        clearTimeout(timeoutRef.current)
+      }
+    }
+  }, [value, delay])
+
+  return shouldRender ? children : fallback
+}

packages/ui-patterns/AssistantChat/AssistantChatForm.tsx

@@ -1,7 +1,7 @@
 'use client'
 
 import { Loader2 } from 'lucide-react'
-import React, { ChangeEvent, useRef } from 'react'
+import React, { ChangeEvent, memo, useRef } from 'react'
 import { ExpandingTextArea } from 'ui'
 import { cn } from 'ui/src/lib/utils'
 
@@ -28,7 +28,7 @@ export interface FormProps extends React.FormHTMLAttributes<HTMLFormElement> {
   placeholder?: string
 }
 
-const AssistantChatForm = React.forwardRef<HTMLFormElement, FormProps>(
+const AssistantChatFormComponent = React.forwardRef<HTMLFormElement, FormProps>(
   (
     {
       loading = false,
@@ -126,6 +126,6 @@ const AssistantChatForm = React.forwardRef<HTMLFormElement, FormProps>(
   }
 )
 
-AssistantChatForm.displayName = 'AssistantChatForm'
+AssistantChatFormComponent.displayName = 'AssistantChatFormComponent'
 
-export { AssistantChatForm }
+export const AssistantChatForm = memo(AssistantChatFormComponent)

packages/ui/src/layout/ai-icon-animation/ai-icon-animation.tsx

@@ -1,7 +1,7 @@
 'use client'
 
 import { motion, useMotionValue, useSpring } from 'framer-motion'
-import { useRef, useState } from 'react'
+import { useRef, useState, memo } from 'react'
 import { cn } from '../../lib/utils'
 
 interface AiIconAnimationProps {
@@ -11,7 +11,7 @@ interface AiIconAnimationProps {
   allowHoverEffect?: boolean
 }
 
-export const AiIconAnimation = ({
+const AiIconAnimationComponent = ({
   size = 24,
   loading = false,
   className,
@@ -122,3 +122,5 @@ export const AiIconAnimation = ({
     </div>
   )
 }
+
+export const AiIconAnimation = memo(AiIconAnimationComponent)