Feat/save assistant messages (supabase#30823)

* save chat messages

* just ai state in subscribe

* fix hydration issue

* check for url param to open assistant

* fix dashboard history

* fix ts

* safeguard against parsing errors

* Fix the type errors.

* Another try.

---------

Co-authored-by: Alaister Young <[email protected]>
Co-authored-by: Ivan Vasilov <[email protected]>

Author: 3 people

apps/studio/components/layouts/ProjectLayout/ProjectLayout.tsx

@@ -1,6 +1,6 @@
 import Head from 'next/head'
 import { useRouter } from 'next/router'
-import { forwardRef, Fragment, PropsWithChildren, ReactNode, useEffect } from 'react'
+import { forwardRef, Fragment, PropsWithChildren, ReactNode, useEffect, useState } from 'react'
 
 import { useParams } from 'common'
 import ProjectAPIDocs from 'components/interfaces/ProjectAPIDocs/ProjectAPIDocs'
@@ -110,6 +110,12 @@ const ProjectLayout = forwardRef<HTMLDivElement, PropsWithChildren<ProjectLayout
       router.pathname === '/project/[ref]' || router.pathname.includes('/project/[ref]/settings')
     const showPausedState = isPaused && !ignorePausedState
 
+    const [isClient, setIsClient] = useState(false)
+
+    useEffect(() => {
+      setIsClient(true)
+    }, [])
+
     useEffect(() => {
       const handler = (e: KeyboardEvent) => {
         if (e.metaKey && e.code === 'KeyI' && !e.altKey && !e.shiftKey) {
@@ -195,7 +201,7 @@ const ProjectLayout = forwardRef<HTMLDivElement, PropsWithChildren<ProjectLayout
                       )}
                     </main>
                   </ResizablePanel>
-                  {aiAssistantPanel.open && (
+                  {isClient && aiAssistantPanel.open && (
                     <>
                       <ResizableHandle />
                       <ResizablePanel

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

@@ -135,7 +135,6 @@ export const AIAssistant = ({
     id,
     api: `${BASE_PATH}/api/ai/sql/generate-v3`,
     maxSteps: 5,
-    // [Joshen] Not currently used atm, but initialMessages will be for...
     initialMessages,
     body: {
       includeSchemaMetadata,
@@ -144,6 +143,11 @@ export const AIAssistant = ({
       schema: currentSchema,
       table: currentTable?.name,
     },
+    onFinish: (message) => {
+      setAiAssistantPanel({
+        messages: [...chatMessages, message],
+      })
+    },
   })
 
   const canUpdateOrganization = useCheckPermissions(PermissionAction.UPDATE, 'organizations')
@@ -185,7 +189,7 @@ export const AIAssistant = ({
       headers: { Authorization: headerData.get('Authorization') ?? '' },
     })
 
-    setAiAssistantPanel({ sqlSnippets: undefined })
+    setAiAssistantPanel({ sqlSnippets: undefined, messages: [...messages, payload] })
     setValue('')
     setAssistantError(undefined)
     setLastSentMessage(payload)
@@ -400,7 +404,7 @@ export const AIAssistant = ({
               </h3>
               {suggestions.title && <p>{suggestions.title}</p>}
               <div className="-mx-3 mt-4 mb-12">
-                {suggestions?.prompts?.map((prompt, idx) => (
+                {suggestions?.prompts?.map((prompt: string, idx: number) => (
                   <Button
                     key={`suggestion-${idx}`}
                     size="small"
@@ -496,7 +500,7 @@ export const AIAssistant = ({
         <div className="p-5 pt-0 z-20 relative">
           {sqlSnippets && sqlSnippets.length > 0 && (
             <div className="mb-2">
-              {sqlSnippets.map((snippet, index) => (
+              {sqlSnippets.map((snippet: string, index: number) => (
                 <CollapsibleCodeBlock
                   key={index}
                   hideLineNumbers
@@ -552,7 +556,9 @@ export const AIAssistant = ({
               event.preventDefault()
               if (includeSchemaMetadata) {
                 const sqlSnippetsString =
-                  sqlSnippets?.map((snippet) => '```sql\n' + snippet + '\n```').join('\n') || ''
+                  sqlSnippets
+                    ?.map((snippet: string) => '```sql\n' + snippet + '\n```')
+                    .join('\n') || ''
                 const valueWithSnippets = [value, sqlSnippetsString].filter(Boolean).join('\n\n')
                 sendMessageToAssistant(valueWithSnippets)
               } else {

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

@@ -1,26 +1,15 @@
+import type { Message as MessageType } from 'ai/react'
 import { uuidv4 } from 'lib/helpers'
-import { useState, useEffect } from 'react'
+import { useState } from 'react'
 import { useAppStateSnapshot } from 'state/app-state'
 import { cn } from 'ui'
 import { AIAssistant } from './AIAssistant'
-import type { Message as MessageType } from 'ai/react'
 
 export const AiAssistantPanel = () => {
   const { aiAssistantPanel, resetAiAssistantPanel } = useAppStateSnapshot()
-  const [initialMessages, setInitialMessages] = useState<MessageType[] | undefined>(undefined)
-
-  useEffect(() => {
-    // set initial state of local messages to the global state if it exists
-    if (aiAssistantPanel.messages) {
-      const messagesCopy = aiAssistantPanel.messages.map((msg) => ({
-        content: msg.content,
-        createdAt: msg.createdAt,
-        role: msg.role,
-        id: msg.id,
-      }))
-      setInitialMessages(messagesCopy)
-    }
-  }, [])
+  const [initialMessages, setInitialMessages] = useState<MessageType[] | undefined>(
+    aiAssistantPanel.messages?.length > 0 ? (aiAssistantPanel.messages as any) : undefined
+  )
 
   const { open } = aiAssistantPanel
   const [chatId, setChatId] = useState(() => uuidv4())

apps/studio/lib/constants/index.ts

@@ -30,6 +30,8 @@ export const USAGE_APPROACHING_THRESHOLD = 0.75
 export const LOCAL_STORAGE_KEYS = {
   RECENTLY_VISITED_ORGANIZATION: 'supabase-organization',
 
+  AI_ASSISTANT_STATE: 'supabase-ai-assistant-state',
+
   UI_PREVIEW_NAVIGATION_LAYOUT: 'supabase-ui-preview-nav-layout',
   UI_PREVIEW_API_SIDE_PANEL: 'supabase-ui-api-side-panel',
   UI_PREVIEW_CLS: 'supabase-ui-cls',

apps/studio/state/app-state.ts

@@ -1,17 +1,8 @@
-import { proxy, snapshot, useSnapshot } from 'valtio'
-
+import { proxy, subscribe, snapshot, useSnapshot } from 'valtio'
+import type { Message as MessageType } from 'ai/react'
 import { SupportedAssistantEntities } from 'components/ui/AIAssistantPanel/AIAssistant.types'
 import { LOCAL_STORAGE_KEYS } from 'lib/constants'
 import { LOCAL_STORAGE_KEYS as COMMON_LOCAL_STORAGE_KEYS } from 'common'
-import type { Message as MessageType } from 'ai/react'
-
-const EMPTY_DASHBOARD_HISTORY: {
-  sql?: string
-  editor?: string
-} = {
-  sql: undefined,
-  editor: undefined,
-}
 
 export type CommonDatabaseEntity = {
   id: number
@@ -27,7 +18,7 @@ export type SuggestionsType = {
 
 type AiAssistantPanelType = {
   open: boolean
-  messages?: MessageType[] | undefined
+  messages: MessageType[]
   initialInput: string
   sqlSnippets?: string[]
   suggestions?: SuggestionsType
@@ -39,9 +30,14 @@ type AiAssistantPanelType = {
   tables: { schema: string; name: string }[]
 }
 
+type DashboardHistoryType = {
+  sql?: string
+  editor?: string
+}
+
 const INITIAL_AI_ASSISTANT: AiAssistantPanelType = {
   open: false,
-  messages: undefined,
+  messages: [],
   sqlSnippets: undefined,
   initialInput: '',
   suggestions: undefined,
@@ -51,9 +47,76 @@ const INITIAL_AI_ASSISTANT: AiAssistantPanelType = {
   tables: [],
 }
 
+const EMPTY_DASHBOARD_HISTORY: DashboardHistoryType = {
+  sql: undefined,
+  editor: undefined,
+}
+
+const getInitialState = () => {
+  if (typeof window === 'undefined') {
+    return {
+      aiAssistantPanel: INITIAL_AI_ASSISTANT,
+      dashboardHistory: EMPTY_DASHBOARD_HISTORY,
+      activeDocsSection: ['introduction'],
+      docsLanguage: 'js',
+      showProjectApiDocs: false,
+      isOptedInTelemetry: false,
+      showEnableBranchingModal: false,
+      showFeaturePreviewModal: false,
+      selectedFeaturePreview: '',
+      showAiSettingsModal: false,
+      showGenerateSqlModal: false,
+      navigationPanelOpen: false,
+      navigationPanelJustClosed: false,
+    }
+  }
+
+  const stored = localStorage.getItem(LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE)
+
+  const urlParams = new URLSearchParams(window.location.search)
+  const aiAssistantPanelOpenParam = urlParams.get('aiAssistantPanelOpen')
+
+  let parsedAiAssistant = INITIAL_AI_ASSISTANT
+
+  try {
+    if (stored) {
+      parsedAiAssistant = JSON.parse(stored, (key, value) => {
+        if (key === 'createdAt' && value) {
+          return new Date(value)
+        }
+        return value
+      })
+    }
+  } catch {
+    // Ignore parsing errors
+  }
+
+  return {
+    aiAssistantPanel: {
+      ...parsedAiAssistant,
+      open:
+        aiAssistantPanelOpenParam !== null
+          ? aiAssistantPanelOpenParam === 'true'
+          : parsedAiAssistant.open,
+    },
+    dashboardHistory: EMPTY_DASHBOARD_HISTORY,
+    activeDocsSection: ['introduction'],
+    docsLanguage: 'js',
+    showProjectApiDocs: false,
+    isOptedInTelemetry: false,
+    showEnableBranchingModal: false,
+    showFeaturePreviewModal: false,
+    selectedFeaturePreview: '',
+    showAiSettingsModal: false,
+    showGenerateSqlModal: false,
+    navigationPanelOpen: false,
+    navigationPanelJustClosed: false,
+  }
+}
+
 export const appState = proxy({
-  // [Joshen] Last visited "entity" for any page that we wanna track
-  dashboardHistory: EMPTY_DASHBOARD_HISTORY,
+  ...getInitialState(),
+
   setDashboardHistory: (ref: string, key: 'sql' | 'editor', id: string) => {
     if (appState.dashboardHistory[key] !== id) {
       appState.dashboardHistory[key] = id
@@ -84,22 +147,27 @@ export const appState = proxy({
       localStorage.setItem(COMMON_LOCAL_STORAGE_KEYS.TELEMETRY_CONSENT, value.toString())
     }
   },
+
   showEnableBranchingModal: false,
   setShowEnableBranchingModal: (value: boolean) => {
     appState.showEnableBranchingModal = value
   },
+
   showFeaturePreviewModal: false,
   setShowFeaturePreviewModal: (value: boolean) => {
     appState.showFeaturePreviewModal = value
   },
+
   selectedFeaturePreview: '',
   setSelectedFeaturePreview: (value: string) => {
     appState.selectedFeaturePreview = value
   },
+
   showAiSettingsModal: false,
   setShowAiSettingsModal: (value: boolean) => {
     appState.showAiSettingsModal = value
   },
+
   showGenerateSqlModal: false,
   setShowGenerateSqlModal: (value: boolean) => {
     appState.showGenerateSqlModal = value
@@ -109,18 +177,14 @@ export const appState = proxy({
   navigationPanelJustClosed: false,
   setNavigationPanelOpen: (value: boolean, trackJustClosed: boolean = false) => {
     if (value === false) {
-      // If closing navigation panel by clicking on icon/button, nav bar should not open again until mouse leaves nav bar
       if (trackJustClosed) {
         appState.navigationPanelOpen = false
         appState.navigationPanelJustClosed = true
       } else {
-        // If closing navigation panel by leaving nav bar, nav bar can open again when mouse re-enter
         appState.navigationPanelOpen = false
         appState.navigationPanelJustClosed = false
       }
     } else {
-      // If opening nav panel, check if it was just closed by a nav icon/button click
-      // If yes, do not open nav panel, otherwise open as per normal
       if (appState.navigationPanelJustClosed === false) {
         appState.navigationPanelOpen = true
       }
@@ -137,7 +201,6 @@ export const appState = proxy({
     }
   },
 
-  aiAssistantPanel: INITIAL_AI_ASSISTANT as AiAssistantPanelType,
   setAiAssistantPanel: (value: Partial<AiAssistantPanelType>) => {
     const hasEntityChanged = value.entity?.id !== appState.aiAssistantPanel.entity?.id
 
@@ -149,6 +212,16 @@ export const appState = proxy({
   },
 })
 
+// Set up localStorage subscription
+if (typeof window !== 'undefined') {
+  subscribe(appState, () => {
+    localStorage.setItem(
+      LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE,
+      JSON.stringify(appState.aiAssistantPanel)
+    )
+  })
+}
+
 export const getAppStateSnapshot = () => snapshot(appState)
 
 export const useAppStateSnapshot = (options?: Parameters<typeof useSnapshot>[1]) =>