diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index 5fd65dea4df..40d159e925e 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/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 = ({ ) }), @@ -313,12 +307,12 @@ export const AIAssistant = ({ )} {hasMessages ? ( - +
{renderedMessages} {(last(messages)?.role === 'user' || last(messages)?.content?.length === 0) && (
- +
Thinking
- +
)}
- +
) : suggestions ? (

diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx index 1eb5ed6410e..509fc5f88b1 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx +++ b/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({ 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 + }) => ( + + Writing SQL... +

+ } + > + + + ) +) +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 ( +
+ {language === 'sql' ? ( + readOnly ? ( + + ) : ( + + ) + ) : ( + code]:m-0 [&>code>span]:flex [&>code>span]:flex-wrap [&>code]:block [&>code>span]:text-foreground' + )} + /> + )} +
+ ) +} + +const OrderedList = memo(({ children }: { children: ReactNode }) => ( +
    {children}
+)) +OrderedList.displayName = 'OrderedList' + +const ListItem = memo(({ children }: { children: ReactNode }) => ( +
  • {children}
  • +)) +ListItem.displayName = 'ListItem' + +const Heading3 = memo(({ children }: { children: ReactNode }) => ( +

    {children}

    +)) +Heading3.displayName = 'Heading3' + +const InlineCode = memo(({ className, children }: { className?: string; children: ReactNode }) => ( + {children} +)) +InlineCode.displayName = 'InlineCode' + +const Link = memo(({ href, children }: { href?: string; children: ReactNode }) => ( + + {children} + +)) +Link.displayName = 'Link' + +const baseMarkdownComponents: Partial = { + 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) { + useEffect(() => { + return () => { + console.log('unmounting parent') + } + }, []) + const isUser = role === 'user' + const allMarkdownComponents = useMemo( + () => ({ + ...markdownComponents, + ...baseMarkdownComponents, + }), + [] + ) + if (!content) return null return ( - - {children} - - {variant === 'warning' && } - - {action} - -
    - {isUser ? ( -
    - -
    - ) : ( - + +
    { - 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' && } - 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 ( -
    - {language === 'sql' ? ( - readOnly ? ( - - ) : ( - - ) - ) : ( - code]:m-0 [&>code>span]:flex [&>code>span]:flex-wrap [&>code]:block [&>code>span]:text-foreground' - )} - /> - )} -
    - ) - }, - ol: (props: any) => { - return
      {props.children}
    - }, - li: (props: any) => { - return
  • {props.children}
  • - }, - h3: (props: any) => { - return

    {props.children}

    - }, - code: (props: any) => { - return {props.children} - }, - a: (props: any) => { - return ( - - {props.children} - - ) - }, - }} - > - {content} - +
    + {isUser ? ( +
    + +
    + ) : ( + + )} + + {content} + +
    - +
    ) } diff --git a/apps/studio/components/ui/DebouncedComponent.tsx b/apps/studio/components/ui/DebouncedComponent.tsx new file mode 100644 index 00000000000..ca3fee910f2 --- /dev/null +++ b/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 =
    Loading...
    , + children, +}: DebouncedComponentProps) { + const [shouldRender, setShouldRender] = useState(false) + const timeoutRef = useRef() + 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 +} diff --git a/packages/ui-patterns/AssistantChat/AssistantChatForm.tsx b/packages/ui-patterns/AssistantChat/AssistantChatForm.tsx index 4e384513e0b..8c609acd9ef 100644 --- a/packages/ui-patterns/AssistantChat/AssistantChatForm.tsx +++ b/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 { placeholder?: string } -const AssistantChatForm = React.forwardRef( +const AssistantChatFormComponent = React.forwardRef( ( { loading = false, @@ -126,6 +126,6 @@ const AssistantChatForm = React.forwardRef( } ) -AssistantChatForm.displayName = 'AssistantChatForm' +AssistantChatFormComponent.displayName = 'AssistantChatFormComponent' -export { AssistantChatForm } +export const AssistantChatForm = memo(AssistantChatFormComponent) diff --git a/packages/ui/src/layout/ai-icon-animation/ai-icon-animation.tsx b/packages/ui/src/layout/ai-icon-animation/ai-icon-animation.tsx index 3f838bada84..eaa77fa4404 100644 --- a/packages/ui/src/layout/ai-icon-animation/ai-icon-animation.tsx +++ b/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 = ({
    ) } + +export const AiIconAnimation = memo(AiIconAnimationComponent)