From c8058fd9fc57706f674d7b1ee472aa40366d9c8b Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Wed, 26 Mar 2025 20:23:50 +0800
Subject: [PATCH] Midway

---
 .../ui/AIAssistantPanel/AIAssistant.tsx       | 79 +++++++++++--------
 .../AIAssistantChatSelector.tsx               |  8 +-
 .../ui/AIAssistantPanel/Message.tsx           | 14 +++-
 .../ui/AIAssistantPanel/MessageMarkdown.tsx   | 16 +++-
 .../components/ui/QueryBlock/QueryBlock.tsx   | 23 +++---
 apps/studio/hooks/useAssistant.ts             | 36 +++++----
 apps/studio/state/app-state.ts                | 10 ++-
 7 files changed, 110 insertions(+), 76 deletions(-)

diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
index 656f1bb78c1..770c3719efa 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
@@ -3,9 +3,9 @@ import type { Message as MessageType } from 'ai/react'
 import { useChat } from 'ai/react'
 import { AnimatePresence, motion } from 'framer-motion'
 import { last } from 'lodash'
-import { ArrowDown, FileText, Info, RefreshCw, X, Plus } from 'lucide-react'
+import { ArrowDown, FileText, Info, Plus, RefreshCw, X } from 'lucide-react'
 import { useRouter } from 'next/router'
-import { memo, useEffect, useMemo, useRef, useState, useCallback } from 'react'
+import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import { toast } from 'sonner'
 
 import { useParams, useSearchParamsShallow } from 'common/hooks'
@@ -24,6 +24,7 @@ import { useOrgOptedIntoAi } from 'hooks/misc/useOrgOptedIntoAi'
 import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization'
 import { useSelectedProject } from 'hooks/misc/useSelectedProject'
 import { useFlag } from 'hooks/ui/useFlag'
+import { useAssistant } from 'hooks/useAssistant'
 import { BASE_PATH, IS_PLATFORM, OPT_IN_TAGS } from 'lib/constants'
 import uuidv4 from 'lib/uuid'
 import { useAppStateSnapshot } from 'state/app-state'
@@ -39,23 +40,33 @@ import {
 } from 'ui'
 import { Admonition, AssistantChatForm, GenericSkeletonLoader } from 'ui-patterns'
 import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
+import { ButtonTooltip } from '../ButtonTooltip'
 import DotGrid from '../DotGrid'
+import { AIAssistantChatSelector } from './AIAssistantChatSelector'
 import AIOnboarding from './AIOnboarding'
 import CollapsibleCodeBlock from './CollapsibleCodeBlock'
 import { Message } from './Message'
 import { useAutoScroll } from './hooks'
-import { AIAssistantChatSelector } from './AIAssistantChatSelector'
-import { useAssistant } from 'hooks/useAssistant'
 
 const MemoizedMessage = memo(
-  ({ message, isLoading }: { message: MessageType; isLoading: boolean }) => {
+  ({
+    message,
+    isLoading,
+    onResults,
+  }: {
+    message: MessageType
+    isLoading: boolean
+    onResults: ({ title, results }: { title?: string; results: any[] }) => void
+  }) => {
     return (
       <Message
         key={message.id}
+        id={message.id}
         role={message.role}
         content={message.content}
         readOnly={message.role === 'user'}
         isLoading={isLoading}
+        onResults={onResults}
       />
     )
   }
@@ -83,6 +94,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   const {
     messages: assistantMessages,
     saveMessage,
+    updateMessage,
     clearMessages,
     closeAssistant,
     clearSqlSnippets,
@@ -180,6 +192,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
             key={message.id}
             message={message}
             isLoading={isChatLoading && message.id === chatMessages[chatMessages.length - 1].id}
+            onResults={(props) => updateMessage({ id: message.id, ...props })}
           />
         )
       }),
@@ -302,36 +315,32 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
               </div>
               <div className="flex gap-2">
                 <AIAssistantChatSelector />
-                <Tooltip>
-                  <TooltipTrigger asChild>
-                    <Button
-                      type="default"
-                      size="tiny"
-                      icon={<Plus size={14} />}
-                      onClick={() => {
-                        if (project?.ref) {
-                          newChat()
-                        }
-                      }}
-                      className="h-7 w-7 p-0"
-                    />
-                  </TooltipTrigger>
-                  <TooltipContent>New chat</TooltipContent>
-                </Tooltip>
-                <Tooltip>
-                  <TooltipTrigger asChild>
-                    <Button
-                      type="default"
-                      size="tiny"
-                      icon={<RefreshCw size={14} />}
-                      onClick={handleClearMessages}
-                      className="h-7 w-7 p-0"
-                      disabled={isChatLoading}
-                    />
-                  </TooltipTrigger>
-                  <TooltipContent>Clear chat messages</TooltipContent>
-                </Tooltip>
-                <Button type="default" className="w-7 h-7" onClick={closeAssistant} icon={<X />} />
+                <ButtonTooltip
+                  type="default"
+                  size="tiny"
+                  icon={<Plus size={14} />}
+                  onClick={() => {
+                    if (project?.ref) newChat()
+                  }}
+                  tooltip={{ content: { side: 'bottom', text: 'New chat' } }}
+                  className="h-7 w-7 p-0"
+                />
+                <ButtonTooltip
+                  type="default"
+                  size="tiny"
+                  icon={<RefreshCw size={14} />}
+                  onClick={handleClearMessages}
+                  className="h-7 w-7 p-0"
+                  tooltip={{ content: { side: 'bottom', text: 'Clear chat messages' } }}
+                  disabled={isChatLoading}
+                />
+                <ButtonTooltip
+                  type="default"
+                  className="w-7 h-7"
+                  onClick={closeAssistant}
+                  icon={<X />}
+                  tooltip={{ content: { side: 'bottom', text: 'Close Assistant' } }}
+                />
               </div>
             </div>
             {!includeSchemaMetadata && selectedOrganization && (
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx
index ad60060f628..d5abf36daaa 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx
@@ -1,7 +1,7 @@
-import { useState } from 'react'
-import { Check, Edit, MessageSquare, Trash, X } from 'lucide-react'
 import { useAssistant } from 'hooks/useAssistant'
+import { Check, Edit, MessageSquare, Trash, X } from 'lucide-react'
 import { useRouter } from 'next/router'
+import { useState } from 'react'
 import {
   Button,
   cn,
@@ -16,6 +16,7 @@ import {
   PopoverContent_Shadcn_,
   PopoverTrigger_Shadcn_,
 } from 'ui'
+import { ButtonTooltip } from '../ButtonTooltip'
 
 interface AIAssistantChatSelectorProps {
   className?: string
@@ -86,11 +87,12 @@ export const AIAssistantChatSelector = ({ className }: AIAssistantChatSelectorPr
   return (
     <Popover_Shadcn_ open={chatSelectorOpen} onOpenChange={setChatSelectorOpen}>
       <PopoverTrigger_Shadcn_ asChild>
-        <Button
+        <ButtonTooltip
           type="default"
           size="tiny"
           icon={<MessageSquare size={14} />}
           className={cn('h-7 w-7', className)}
+          tooltip={{ content: { side: 'bottom', text: 'View chats' } }}
         />
       </PopoverTrigger_Shadcn_>
       <PopoverContent_Shadcn_ className="w-[250px] p-0" align="end">
diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx
index 4ae075b24ca..b9fe6839799 100644
--- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/Message.tsx
@@ -19,19 +19,21 @@ const baseMarkdownComponents: Partial<Components> = {
   h3: Heading3,
   code: InlineCode,
   a: Link,
-  pre: MarkdownPre,
 }
 
 interface MessageProps {
+  id: string
   role: 'function' | 'system' | 'user' | 'assistant' | 'data' | 'tool'
   content?: string
   isLoading: boolean
   readOnly?: boolean
   action?: React.ReactNode
   variant?: 'default' | 'warning'
+  onResults: ({ title, results }: { title?: string; results: any[] }) => void
 }
 
 export const Message = function Message({
+  id,
   role,
   content,
   isLoading,
@@ -39,10 +41,16 @@ export const Message = function Message({
   children,
   action = null,
   variant = 'default',
+  onResults,
 }: PropsWithChildren<MessageProps>) {
   const isUser = role === 'user'
-  const allMarkdownComponents = useMemo(
-    () => ({ ...markdownComponents, ...baseMarkdownComponents }),
+
+  const allMarkdownComponents: Partial<Components> = useMemo(
+    () => ({
+      ...markdownComponents,
+      ...baseMarkdownComponents,
+      pre: ({ children }) => <MarkdownPre onResults={onResults}>{children}</MarkdownPre>,
+    }),
     []
   )
 
diff --git a/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx b/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx
index 1a9375afe1d..8232577af91 100644
--- a/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx
@@ -10,12 +10,12 @@ import { useProfile } from 'lib/profile'
 import { Dashboards } from 'types'
 import { Badge, cn, CodeBlock, CodeBlockLang } from 'ui'
 import { DebouncedComponent } from '../DebouncedComponent'
+import { EdgeFunctionBlock } from '../EdgeFunctionBlock/EdgeFunctionBlock'
 import { QueryBlock } from '../QueryBlock/QueryBlock'
 import { AssistantSnippetProps } from './AIAssistant.types'
 import { identifyQueryType } from './AIAssistant.utils'
 import CollapsibleCodeBlock from './CollapsibleCodeBlock'
 import { MessageContext } from './Message'
-import { EdgeFunctionBlock } from '../EdgeFunctionBlock/EdgeFunctionBlock'
 
 export const OrderedList = memo(({ children }: { children: ReactNode }) => (
   <ol className="flex flex-col gap-y-4">{children}</ol>
@@ -62,6 +62,7 @@ const MemoizedQueryBlock = memo(
     isDraggable,
     runQuery,
     onRunQuery,
+    onResults,
     onDragStart,
     onUpdateChartConfig,
   }: {
@@ -74,6 +75,7 @@ const MemoizedQueryBlock = memo(
     isDraggable: boolean
     runQuery: boolean
     onRunQuery: (queryType: 'select' | 'mutation') => void
+    onResults: (results: any[]) => void
     onDragStart: (e: DragEvent<Element>) => void
     onUpdateChartConfig?: ({
       chart,
@@ -84,7 +86,7 @@ const MemoizedQueryBlock = memo(
     }) => void
   }) => (
     <DebouncedComponent
-      delay={500}
+      delay={isLoading ? 500 : 0}
       value={sql}
       fallback={
         <div className="bg-surface-100 border-overlay rounded border shadow-sm px-3 py-2 text-xs">
@@ -119,6 +121,7 @@ const MemoizedQueryBlock = memo(
         draggable={isDraggable}
         runQuery={runQuery}
         onRunQuery={onRunQuery}
+        onResults={onResults}
         onDragStart={onDragStart}
         onUpdateChartConfig={onUpdateChartConfig}
       />
@@ -127,7 +130,13 @@ const MemoizedQueryBlock = memo(
 )
 MemoizedQueryBlock.displayName = 'MemoizedQueryBlock'
 
-export const MarkdownPre = ({ children }: { children: any }) => {
+export const MarkdownPre = ({
+  onResults,
+  children,
+}: {
+  onResults: ({ title, results }: { title?: string; results: any[] }) => void
+  children: any
+}) => {
   const router = useRouter()
   const { profile } = useProfile()
   const { isLoading, readOnly } = useContext(MessageContext)
@@ -213,6 +222,7 @@ export const MarkdownPre = ({ children }: { children: any }) => {
             isDraggable={isDraggableToReports}
             runQuery={runQuery}
             onRunQuery={onRunQuery}
+            onResults={(results) => onResults({ title: snippetProps.title, results })}
             onUpdateChartConfig={({ chartConfig: config }) => {
               chartConfig.current = { ...chartConfig.current, ...config }
             }}
diff --git a/apps/studio/components/ui/QueryBlock/QueryBlock.tsx b/apps/studio/components/ui/QueryBlock/QueryBlock.tsx
index 3ce07d91e18..d9a229d0c78 100644
--- a/apps/studio/components/ui/QueryBlock/QueryBlock.tsx
+++ b/apps/studio/components/ui/QueryBlock/QueryBlock.tsx
@@ -11,24 +11,15 @@ import { useProjectContext } from 'components/layouts/ProjectLayout/ProjectConte
 import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
 import { Parameter, parseParameters } from 'lib/sql-parameters'
 import { Dashboards } from 'types'
-import {
-  Button,
-  ChartContainer,
-  ChartTooltip,
-  ChartTooltipContent,
-  cn,
-  CodeBlock,
-  SQL_ICON,
-} from 'ui'
-import { Admonition } from 'ui-patterns'
+import { ChartContainer, ChartTooltip, ChartTooltipContent, cn, CodeBlock, SQL_ICON } from 'ui'
 import ShimmeringLoader from 'ui-patterns/ShimmeringLoader'
 import { containsUnknownFunction, isReadOnlySelect } from '../AIAssistantPanel/AIAssistant.utils'
 import { ButtonTooltip } from '../ButtonTooltip'
+import SqlWarningAdmonition from '../SqlWarningAdmonition'
 import { BlockViewConfiguration } from './BlockViewConfiguration'
 import { EditQueryButton } from './EditQueryButton'
 import { ParametersPopover } from './ParametersPopover'
 import { getCumulativeResults } from './QueryBlock.utils'
-import SqlWarningAdmonition from '../SqlWarningAdmonition'
 
 export const DEFAULT_CHART_CONFIG: ChartConfig = {
   type: 'bar',
@@ -71,10 +62,12 @@ interface QueryBlockProps {
   tooltip?: ReactNode
   /** Not implemented yet: Will be the next part of ReportsV2 */
   onSetParameter?: (params: Parameter[]) => void
-  /** Optional callback the SQL query is run */
+  /** Optional callback when the SQL query is run */
   onRunQuery?: (queryType: 'select' | 'mutation') => void
   /** Optional callback on drag start */
   onDragStart?: (e: DragEvent<Element>) => void
+  /** Optional: callback when the results are returned from running the SQL query*/
+  onResults?: (results: any[]) => void
 
   // [Joshen] Params below are currently only used by ReportsV2 (Might revisit to see how to improve these)
   /** Optional height set to render the SQL query (Used in Reports) */
@@ -117,6 +110,7 @@ export const QueryBlock = ({
   noResultPlaceholder = null,
   tooltip,
   onRunQuery,
+  onResults,
   onSetParameter,
   onUpdateChartConfig,
   onDragStart,
@@ -141,7 +135,10 @@ export const QueryBlock = ({
   const isReadOnlySelectSQL = isReadOnlySelect(sql ?? '')
 
   const { mutate: execute, isLoading: isExecuting } = useExecuteSqlMutation({
-    onSuccess: (data) => setQueryResult(data.result),
+    onSuccess: (data) => {
+      onResults?.(data.result)
+      setQueryResult(data.result)
+    },
   })
 
   const handleExecute = () => {
diff --git a/apps/studio/hooks/useAssistant.ts b/apps/studio/hooks/useAssistant.ts
index a32fb14a87f..edc35baec63 100644
--- a/apps/studio/hooks/useAssistant.ts
+++ b/apps/studio/hooks/useAssistant.ts
@@ -1,7 +1,6 @@
 import type { Message, Message as MessageType } from 'ai/react'
 import { useCallback, useEffect, useMemo } from 'react'
-import { useAppStateSnapshot } from 'state/app-state'
-import { ChatSession } from 'state/app-state'
+import { ChatSession, useAppStateSnapshot } from 'state/app-state'
 
 export interface UseAssistantOptions {
   projectRef?: string
@@ -69,7 +68,7 @@ export function useAssistant(options?: UseAssistantOptions) {
 
           setAiAssistantPanel({
             chats: {
-              ...(aiAssistantPanel.chats || {}),
+              ...(chats || {}),
               [chatId]: newChat,
             } as any,
             activeChatId: chatId,
@@ -93,7 +92,7 @@ export function useAssistant(options?: UseAssistantOptions) {
 
       setAiAssistantPanel({
         chats: {
-          ...(aiAssistantPanel.chats || {}),
+          ...(chats || {}),
           [chatId]: newChat,
         } as any,
         activeChatId: chatId,
@@ -102,7 +101,7 @@ export function useAssistant(options?: UseAssistantOptions) {
       return chatId
     }
     return undefined
-  }, [projectRef, setAiAssistantPanel, aiAssistantPanel.chats])
+  }, [projectRef, setAiAssistantPanel, chats])
 
   const handleSelectChat = useCallback(
     (id: string) => {
@@ -154,16 +153,10 @@ export function useAssistant(options?: UseAssistantOptions) {
         // Create new chats object with the updated chat
         const updatedChats = {
           ...chats,
-          [id]: {
-            ...chat,
-            name,
-            updatedAt: new Date(),
-          },
-        } as Record<string, ChatSession>
+          [id]: { ...chat, name, updatedAt: new Date() },
+        }
 
-        setAiAssistantPanel({
-          chats: updatedChats,
-        })
+        setAiAssistantPanel({ chats: updatedChats as any })
       }
     },
     [projectChatsRecord]
@@ -213,7 +206,7 @@ export function useAssistant(options?: UseAssistantOptions) {
             },
           }
           setAiAssistantPanel({
-            chats: chatUpdate.chats as Record<string, ChatSession>,
+            chats: chatUpdate.chats as any,
           })
         }
       }
@@ -221,6 +214,18 @@ export function useAssistant(options?: UseAssistantOptions) {
     [activeChatId, projectChatsRecord]
   )
 
+  const handleUpdateMessage = useCallback(
+    ({ id, title, results }: { id: string; title?: string; results: any[] }) => {
+      if (projectRef && currentChatBelongsToProject && activeChatId) {
+        // Get the existing messages for the active chat
+        const existingMessages: readonly MessageType[] =
+          projectChatsRecord[activeChatId]?.messages || []
+        console.log({ existingMessages })
+      }
+    },
+    [activeChatId, projectChatsRecord]
+  )
+
   const openAssistant = useCallback(
     (options?: Partial<typeof aiAssistantPanel>) => {
       // If opening the assistant, ensure we have a valid project chat
@@ -294,5 +299,6 @@ export function useAssistant(options?: UseAssistantOptions) {
     clearSqlSnippets,
     clearSuggestions,
     setSqlSnippets,
+    updateMessage: handleUpdateMessage,
   }
 }
diff --git a/apps/studio/state/app-state.ts b/apps/studio/state/app-state.ts
index 364f5aa3f41..d22109f31d5 100644
--- a/apps/studio/state/app-state.ts
+++ b/apps/studio/state/app-state.ts
@@ -1,9 +1,9 @@
 import type { Message as MessageType } from 'ai/react'
 import { LOCAL_STORAGE_KEYS as COMMON_LOCAL_STORAGE_KEYS } from 'common'
-import { LOCAL_STORAGE_KEYS } from 'lib/constants'
-import { SupportedAssistantEntities } from 'components/ui/AIAssistantPanel/AIAssistant.types'
-import { proxy, snapshot, subscribe, useSnapshot } from 'valtio'
 import { SQL_TEMPLATES } from 'components/interfaces/SQLEditor/SQLEditor.queries'
+import { SupportedAssistantEntities } from 'components/ui/AIAssistantPanel/AIAssistant.types'
+import { LOCAL_STORAGE_KEYS } from 'lib/constants'
+import { proxy, snapshot, subscribe, useSnapshot } from 'valtio'
 
 export type CommonDatabaseEntity = {
   id: number
@@ -23,11 +23,13 @@ export type Template = {
   content: string
 }
 
+type AssistantMessageType = MessageType & { results?: any[] }
+
 export type ChatSession = {
   id: string
   name: string
   projectRef: string
-  messages: readonly MessageType[]
+  messages: readonly AssistantMessageType[]
   createdAt: Date
   updatedAt: Date
 }