fixes and refinements

apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx

@@ -125,7 +125,11 @@ const FunctionList = ({
                         <DropdownMenuItem
                           className="space-x-2"
                           onClick={() => {
-                            setAiAssistantPanel({ open: true, editor: 'functions', entity: x })
+                            setAiAssistantPanel({
+                              open: true,
+                              initialInput: 'Help me update this function to do...',
+                              sqlSnippets: [x.complete_statement],
+                            })
                           }}
                         >
                           <Edit size={14} />

apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionsList.tsx

@@ -155,8 +155,7 @@ const FunctionsList = ({
                       onClick={() =>
                         setAiAssistantPanel({
                           open: true,
-                          editor: 'functions',
-                          entity: undefined,
+                          initialInput: `Help me create a new function for the schema ${selectedSchema} that does ...`,
                         })
                       }
                       tooltip={{

apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx

@@ -181,7 +181,7 @@ const SQLEditor = () => {
           const editor = editorRef.current
           const monaco = monacoRef.current
 
-          const startLineNumber = hasSelection ? editor?.getSelection()?.startLineNumber ?? 0 : 0
+          const startLineNumber = hasSelection ? (editor?.getSelection()?.startLineNumber ?? 0) : 0
 
           const formattedError = error.formattedError ?? ''
           const lineError = formattedError.slice(formattedError.indexOf('LINE'))
@@ -237,7 +237,7 @@ const SQLEditor = () => {
       const selection = editor.getSelection()
       const selectedValue = selection ? editor.getModel()?.getValueInRange(selection) : undefined
       const sql = snippet
-        ? (selectedValue || editorRef.current?.getValue()) ?? snippet.snippet.content.sql
+        ? ((selectedValue || editorRef.current?.getValue()) ?? snippet.snippet.content.sql)
         : selectedValue || editorRef.current?.getValue()
       formatQuery(
         {
@@ -277,7 +277,7 @@ const SQLEditor = () => {
         const selectedValue = selection ? editor.getModel()?.getValueInRange(selection) : undefined
 
         const sql = snippet
-          ? (selectedValue || editorRef.current?.getValue()) ?? snippet.snippet.content.sql
+          ? ((selectedValue || editorRef.current?.getValue()) ?? snippet.snippet.content.sql)
           : selectedValue || editorRef.current?.getValue()
 
         let queryHasIssues = false
@@ -406,26 +406,11 @@ const SQLEditor = () => {
     try {
       const snippet = snapV2.snippets[id]
       const result = snapV2.results[id]?.[0]
-
-      const { solution, sql } = await debugSql({
-        sql: snippet.snippet.content.sql.replace(sqlAiDisclaimerComment, '').trim(),
-        errorMessage: result.error.message,
-        entityDefinitions,
+      appSnap.setAiAssistantPanel({
+        open: true,
+        sqlSnippets: [snippet.snippet.content.sql.replace(sqlAiDisclaimerComment, '').trim()],
+        initialInput: `Help me to debug the attached sql snippet which gives the following error: \n\n${result.error.message}`,
       })
-
-      const formattedSql =
-        sqlAiDisclaimerComment +
-        '\n\n' +
-        format(sql, {
-          language: 'postgresql',
-          keywordCase: 'lower',
-        })
-      setDebugSolution(solution)
-      setSourceSqlDiff({
-        original: snippet.snippet.content.sql,
-        modified: formattedSql,
-      })
-      setSelectedDiffType(DiffType.Modification)
     } catch (error: unknown) {
       // [Joshen] There's a tendency for the SQL debug to chuck a lengthy error message
       // that's not relevant for the user - so we prettify it here by avoiding to return the
@@ -719,7 +704,27 @@ const SQLEditor = () => {
                   <motion.button
                     layoutId="ask-ai-input-icon"
                     transition={{ duration: 0.1 }}
-                    onClick={() => aiPanelRef.current?.expand()}
+                    onClick={() => {
+                      const state = getSqlEditorV2StateSnapshot()
+                      const snippet = state.snippets[id]
+                      const editor = editorRef.current
+                      const selection = editor.getSelection()
+                      const selectedValue = selection
+                        ? editor.getModel()?.getValueInRange(selection)
+                        : undefined
+                      const sql = snippet
+                        ? ((selectedValue || editorRef.current?.getValue()) ??
+                          snippet.snippet.content.sql)
+                        : selectedValue || editorRef.current?.getValue()
+
+                      appSnap.setAiAssistantPanel({
+                        open: true,
+                        sqlSnippets: sql ? [sql] : [],
+                        initialInput: sql
+                          ? `Help me make a change to the attached sql snippet`
+                          : '',
+                      })
+                    }}
                     className={cn(
                       'group absolute z-10 rounded-lg right-[24px] top-4 transition-all duration-200 ease-out'
                     )}

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

@@ -1,452 +1,33 @@
-import { useQueryClient } from '@tanstack/react-query'
-import { uniqBy } from 'lodash'
-import { ChevronLeft, Command, CornerDownLeft, Loader2, X } from 'lucide-react'
-import { useRouter } from 'next/router'
-import { useCallback, useEffect, useRef, useState } from 'react'
-import { toast } from 'sonner'
-
-import { useChat } from 'ai/react'
-import { useParams } from 'common'
-import { useIsDatabaseFunctionsAssistantEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext'
-import { generateThreadMessage } from 'components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyEditorPanel.utils'
-import { MessageWithDebug } from 'components/interfaces/SQLEditor/AiAssistantPanel'
-import { sqlAiDisclaimerComment } from 'components/interfaces/SQLEditor/SQLEditor.constants'
-import { DiffType, IStandaloneCodeEditor } from 'components/interfaces/SQLEditor/SQLEditor.types'
-import { suffixWithLimit } from 'components/interfaces/SQLEditor/SQLEditor.utils'
-import Results from 'components/interfaces/SQLEditor/UtilityPanel/Results'
-import { useSqlDebugMutation } from 'data/ai/sql-debug-mutation'
-import { databasePoliciesKeys } from 'data/database-policies/keys'
-import { QueryResponseError, useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
-import { sqlKeys } from 'data/sql/keys'
-import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
-import { usePrevious } from 'hooks/deprecated'
-import { useLocalStorage } from 'hooks/misc/useLocalStorage'
-import { useSelectedProject } from 'hooks/misc/useSelectedProject'
-import { SqlEditor } from 'icons'
-import {
-  BASE_PATH,
-  LOCAL_STORAGE_KEYS,
-  TELEMETRY_ACTIONS,
-  TELEMETRY_CATEGORIES,
-  TELEMETRY_LABELS,
-} from 'lib/constants'
-import { detectOS, uuidv4 } from 'lib/helpers'
+import { useRef, useState } from 'react'
+import { IStandaloneCodeEditor } from 'components/interfaces/SQLEditor/SQLEditor.types'
+import { uuidv4 } from 'lib/helpers'
 import { useAppStateSnapshot } from 'state/app-state'
-import {
-  Button,
-  cn,
-  Sheet,
-  SheetContent,
-  SheetFooter,
-  SheetHeader,
-  Tooltip_Shadcn_,
-  TooltipContent_Shadcn_,
-  TooltipTrigger_Shadcn_,
-} from 'ui'
-import { Admonition } from 'ui-patterns'
-import CodeEditor from '../CodeEditor/CodeEditor'
+import { cn } from 'ui'
 import { AIAssistant } from './AIAssistant'
-import { generateCTA, generatePlaceholder, generateTitle, validateQuery } from './AIAssistant.utils'
-import { ASSISTANT_SUPPORT_ENTITIES } from './AiAssistant.constants'
-import { useEntityDefinitionQuery } from 'data/database/entity-definition-query'
 
 export const AiAssistantPanel = () => {
-  const os = detectOS()
-  const router = useRouter()
-  const { ref } = useParams()
-  const project = useSelectedProject()
-  const queryClient = useQueryClient()
   const { aiAssistantPanel, setAiAssistantPanel } = useAppStateSnapshot()
-  const isEnabled = useIsDatabaseFunctionsAssistantEnabled()
 
-  const { open, editor, content, entity } = aiAssistantPanel
-  const previousEditor = usePrevious(editor)
-
-  const [isAcknowledged, setIsAcknowledged] = useLocalStorage(
-    LOCAL_STORAGE_KEYS.SQL_SCRATCH_PAD_BANNER_ACKNOWLEDGED,
-    false
-  )
+  const { open, editor, content, entity, initialMessages, initialInput, sqlSnippets } =
+    aiAssistantPanel
 
   const [chatId, setChatId] = useState(() => uuidv4())
-  const [error, setError] = useState<QueryResponseError>()
-  const [results, setResults] = useState<undefined | any[]>(undefined)
-  const [showResults, setShowResults] = useState(false)
-  const [showWarning, setShowWarning] = useState<boolean>(false)
-  const [debugThread, setDebugThread] = useState<MessageWithDebug[]>([])
-
-  const { data: existingDefinition } = useEntityDefinitionQuery({
-    id: entity?.id,
-    type: editor,
-    projectRef: project?.ref,
-    connectionString: project?.connectionString,
-  })
 
   // [Joshen] JFYI I'm opting to just have the assistant always show and not toggle-able
   // I don't really see any negatives of keeping it open (or benefits from hiding) tbh
   // const [showAssistant, setShowAssistant] = useState(true)
-  const title = generateTitle(editor, entity)
-  const placeholder = generatePlaceholder(editor, entity, existingDefinition)
-  const ctaText = generateCTA(editor)
   const editorRef = useRef<IStandaloneCodeEditor | undefined>()
-  const editorModel = editorRef.current?.getModel()
-
-  const numResults = (results ?? []).length
-  const [errorHeader, ...errorContent] =
-    (error?.formattedError?.split('\n') ?? [])?.filter((x: string) => x.length > 0) ?? []
-  const entityContext = ASSISTANT_SUPPORT_ENTITIES.find((x) => x.id === editor)
-
-  const { mutate: sendEvent } = useSendEventMutation()
-  const sendTelemetryEvent = (action: string) => {
-    sendEvent({
-      action,
-      category: TELEMETRY_CATEGORIES.AI_ASSISTANT,
-      label: TELEMETRY_LABELS.QUICK_SQL_EDITOR,
-    })
-  }
-
-  const { mutate: executeSql, isLoading: isExecuting } = useExecuteSqlMutation({
-    onSuccess: async (res) => {
-      // [Joshen] If in a specific editor context mode, assume that intent was to create/update
-      // a database entity - so close it once success. Otherwise it's in Quick SQL mode and we
-      // show the results. Currently though it assumes we're "creating", thinking need to support "updating" too
-      if (editor !== null) {
-        switch (editor) {
-          case 'functions':
-            await queryClient.invalidateQueries(sqlKeys.query(ref, ['functions-list']))
-            break
-          case 'rls-policies':
-            await queryClient.invalidateQueries(databasePoliciesKeys.list(ref))
-            break
-        }
-
-        toast.success(`Successfully created ${entityContext?.name}!`)
-        setAiAssistantPanel({ open: false })
-      } else {
-        setShowResults(true)
-        setResults(res.result)
-      }
-    },
-    onError: (error) => {
-      setError(error)
-      setResults([])
-    },
-  })
-
-  const { append } = useChat({
-    id: chatId,
-    api: `${BASE_PATH}/api/ai/sql/generate-v2`,
-    // [Joshen] Don't need entity definitions if calling here cause support action
-    // via AI here is just to explain code segment, no need for entity definitions context
-    body: {},
-  })
-
-  const { mutateAsync: debugSql } = useSqlDebugMutation({ onError: () => {} })
-
-  const [view, setView] = useState<'assistant' | 'editor'>('assistant')
-
-  const isEditor = view === 'editor'
-
-  const onExecuteSql = (skipValidation = false) => {
-    setError(undefined)
-    setShowWarning(false)
-
-    const query = editorRef.current?.getValue() ?? ''
-    if (query.length === 0) return
-
-    if (editor !== undefined && !skipValidation) {
-      // [Joshen] Some basic validation logic
-      const validated = validateQuery(editor, query)
-      if (!validated) {
-        return setShowWarning(true)
-      }
-    }
-
-    executeSql({
-      sql: suffixWithLimit(query, 100),
-      projectRef: project?.ref,
-      connectionString: project?.connectionString,
-      handleError: (error) => {
-        throw error
-      },
-    })
-  }
-
-  const onFixWithAssistant = async () => {
-    if (!error) return
-
-    const messageId = uuidv4()
-    const query = editorRef.current?.getValue() ?? ''
-
-    const assistantMessageBefore = generateThreadMessage({
-      id: messageId,
-      content: 'Thinking...',
-      isDebug: true,
-    })
-
-    setDebugThread([...debugThread, assistantMessageBefore])
-    sendTelemetryEvent(TELEMETRY_ACTIONS.FIX_WITH_ASSISTANT)
-
-    try {
-      const { solution, sql } = await debugSql({
-        sql: query,
-        errorMessage: error.message,
-      })
-
-      const assistantMessageAfter = generateThreadMessage({
-        id: messageId,
-        content: `${solution}\n\`\`\`sql\n${sql}\n\`\`\``,
-        isDebug: true,
-      })
-      const cleanedMessages = uniqBy([...debugThread, assistantMessageAfter], (m) => m.id)
-      setDebugThread(cleanedMessages)
-    } catch (error) {
-      const assistantMessageAfter = generateThreadMessage({
-        id: messageId,
-        content: `Hmm, Sorry but I'm unable to find a solution for the error that you're facing`,
-        isDebug: true,
-      })
-      const cleanedMessages = uniqBy([...debugThread, assistantMessageAfter], (m) => m.id)
-      setDebugThread(cleanedMessages)
-    }
-  }
-
-  const onExplainSql = (value: string) => {
-    append({
-      role: 'user',
-      createdAt: new Date(),
-      content: `
-        Can you explain this section to me in more detail?\n
-        ${value}
-    `.trim(),
-    })
-    sendTelemetryEvent(TELEMETRY_ACTIONS.EXPLAIN_CODE)
-  }
-
-  const updateEditorWithCheckForDiff = useCallback(
-    ({ id, diffType, sql }: { id: string; diffType: DiffType; sql: string }) => {
-      if (!editorModel) return
-
-      const existingValue = editorRef.current?.getValue() ?? ''
-      if (existingValue.length === 0) {
-        // if the editor is empty, just copy over the code
-        editorRef.current?.executeEdits('apply-ai-message', [
-          {
-            text: `${sqlAiDisclaimerComment}\n\n${sql}`,
-            range: editorModel.getFullModelRange(),
-          },
-        ])
-      } else {
-        const currentSql = editorRef.current?.getValue()
-        editorRef.current?.executeEdits('apply-ai-message', [
-          {
-            text: `${currentSql}\n\n${sqlAiDisclaimerComment}\n\n${sql}`,
-            range: editorModel.getFullModelRange(),
-          },
-        ])
-      }
-      setView('editor')
-    },
-    [editorModel]
-  )
-
-  const onTogglePanel = () => {
-    if (open) {
-      // [Joshen] Save the code content when closing - should allow users to continue from
-      // where they left off IF they are opening the assistant again in the same "editor" context
-      const existingValue = editorRef.current?.getValue() ?? ''
-      setAiAssistantPanel({ open: false, content: existingValue })
-    } else {
-      setAiAssistantPanel({ open: true, editor: null, entity: undefined })
-    }
-  }
-
-  useEffect(() => {
-    if (open) {
-      // [Joshen] Only reset the assistant if the editor changed
-      if (previousEditor !== editor) {
-        setChatId(uuidv4())
-        setError(undefined)
-        setShowWarning(false)
-        setDebugThread([])
-        setResults(undefined)
-        setShowResults(false)
-      } else {
-        // [Joshen] Repopulate the code editor with the content from where the user left off
-        // setTimeout is just to give time for the editorRef to get associated
-        setTimeout(() => {
-          const editorModel = editorRef.current?.getModel()
-          if (content && editorModel) {
-            editorRef.current?.executeEdits('apply-ai-message', [
-              { text: content, range: editorModel.getFullModelRange() },
-            ])
-          }
-        }, 100)
-      }
-    }
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [open])
-
-  // [Joshen] Just to test the concept of a universal assistant of sorts
-  useEffect(() => {
-    const handler = (e: KeyboardEvent) => {
-      if (e.metaKey && e.code === 'KeyI') onTogglePanel()
-    }
-    if (project !== undefined && isEnabled) window.addEventListener('keydown', handler)
-    return () => window.removeEventListener('keydown', handler)
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [project, isEnabled, open])
-
-  // [Joshen] Whenever the deps change recalculate the height of the editor
-  useEffect(() => {
-    editorRef.current?.layout({ width: 0, height: 0 })
-    window.requestAnimationFrame(() => editorRef.current?.layout())
-  }, [error, showWarning, isExecuting, numResults, showResults])
 
   return !open ? null : (
     <div className="w-[400px] xl:relative xl:top-0 absolute right-0 top-[48px] bottom-0 shrink-0 border-l bg">
       <AIAssistant
         id={chatId}
-        className={cn('w-full h-full', isEditor && 'hidden')}
-        debugThread={debugThread}
-        onDiff={updateEditorWithCheckForDiff}
+        initialMessages={initialMessages}
+        sqlSnippets={sqlSnippets}
+        initialInput={initialInput}
+        className={cn('w-full h-full')}
         onResetConversation={() => setChatId(uuidv4())}
       />
-      <div className={cn('flex-1 flex flex-col grow w-full h-full', !isEditor && 'hidden')}>
-        <div className="flex items-center  gap-4 p-4 border-b">
-          <Button
-            type="default"
-            icon={<ChevronLeft strokeWidth={1.5} size={16} />}
-            onClick={() => setView('assistant')}
-          />
-
-          <span className="flex-1 text-sm">{title}</span>
-          <Tooltip_Shadcn_>
-            <TooltipTrigger_Shadcn_ asChild>
-              <Button
-                type="outline"
-                icon={<SqlEditor />}
-                className="h-[24px] w-[24px] px-1"
-                onClick={() => {
-                  const content = editorRef.current?.getValue() ?? ''
-                  router.push(`/project/${ref}/sql/new?content=${encodeURIComponent(content)}`)
-                  setAiAssistantPanel({ open: false })
-                }}
-              />
-            </TooltipTrigger_Shadcn_>
-            <TooltipContent_Shadcn_ side="bottom">Open in SQL Editor</TooltipContent_Shadcn_>
-          </Tooltip_Shadcn_>
-        </div>
-
-        {editor === null && !isAcknowledged && (
-          <Admonition
-            showIcon={false}
-            type="default"
-            title="This is a quick access SQL editor to run queries on your database"
-            className="relative m-0 rounded-none border-x-0 border-t-0 [&>div]:m-0"
-          >
-            <span>
-              Queries written here will not be saved and results are limited up to only 100 rows
-            </span>
-            <Button
-              type="text"
-              icon={<X />}
-              className="px-1.5 absolute top-2 right-2"
-              onClick={() => setIsAcknowledged(true)}
-            />
-          </Admonition>
-        )}
-
-        <div className="flex flex-col h-full justify-between">
-          <div className="relative flex-grow block">
-            <CodeEditor
-              id="assistant-code-editor"
-              language="pgsql"
-              editorRef={editorRef}
-              placeholder={placeholder}
-              actions={{
-                runQuery: { enabled: true, callback: () => onExecuteSql() },
-                explainCode: { enabled: true, callback: onExplainSql },
-                closeAssistant: { enabled: true, callback: () => onTogglePanel() },
-              }}
-            />
-          </div>
-          <div className="flex flex-col">
-            {error !== undefined && (
-              <Admonition
-                type="warning"
-                className="m-0 rounded-none border-x-0 border-b-0 [&>div>div>pre]:text-sm [&>div]:flex [&>div]:flex-col [&>div]:gap-y-2"
-                title={errorHeader || 'Error running SQL query'}
-                description={
-                  <>
-                    <div>
-                      {errorContent.length > 0 ? (
-                        errorContent.map((errorText: string, i: number) => (
-                          <pre key={`err-${i}`} className="font-mono text-xs whitespace-pre-wrap">
-                            {errorText}
-                          </pre>
-                        ))
-                      ) : (
-                        <p className="font-mono text-xs">{error.error}</p>
-                      )}
-                    </div>
-                    <Button type="default" className="w-min" onClick={() => onFixWithAssistant()}>
-                      Fix with Assistant
-                    </Button>
-                  </>
-                }
-              />
-            )}
-            {showWarning && (
-              <Admonition
-                type="default"
-                className="m-0 rounded-none border-x-0 border-b-0 [&>div>pre]:text-sm [&>div]:flex [&>div]:flex-col [&>div]:gap-y-2"
-                title={`Your query doesn't seem to be relevant to ${
-                  entityContext?.id === 'rls-policies'
-                    ? entityContext?.label
-                    : `Database ${entityContext?.label}`
-                }`}
-                description={
-                  <>
-                    <p>Are you sure you want to run this query?</p>
-                    <Button type="default" className="w-min" onClick={() => onExecuteSql(true)}>
-                      Confirm run query
-                    </Button>
-                  </>
-                }
-              />
-            )}
-            {results !== undefined && results.length > 0 && (
-              <>
-                <div className={cn(showResults ? 'h-72 border-t' : 'h-0')}>
-                  <Results rows={results} />
-                </div>
-                <div className="flex items-center justify-between border-t bg-surface-100 py-2 pl-2 pr-5">
-                  <p className="text-xs text-foreground-light">
-                    {results.length} rows
-                    {results.length >= 100 && ` (Limited to only 100 rows)`}
-                  </p>
-                  <Button size="tiny" type="default" onClick={() => setShowResults(!showResults)}>
-                    {showResults ? 'Hide' : 'Show'} results
-                  </Button>
-                </div>
-              </>
-            )}
-            {results !== undefined && results.length === 0 && (
-              <div className="flex items-center justify-between border-t bg-surface-100 h-[43px] pl-2 pr-5">
-                <p className="text-xs text-foreground-light">Success. No rows returned.</p>
-              </div>
-            )}
-            <div className="bg-surface-100 flex items-center !justify-end px-5 py-4 w-full border-t gap-2">
-              <Button type="default" disabled={isExecuting} onClick={() => onTogglePanel()}>
-                Cancel
-              </Button>
-              <Button loading={isExecuting} onClick={() => onExecuteSql()}>
-                {ctaText}
-              </Button>
-            </div>
-          </div>
-        </div>
-      </div>
     </div>
   )
 }

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

@@ -0,0 +1,61 @@
+import { ChevronDown, ChevronUp, X } from 'lucide-react'
+import { useState } from 'react'
+import { Button, CodeBlock, CodeBlockProps, cn } from 'ui'
+
+interface CollapsibleCodeBlockProps extends CodeBlockProps {
+  onRemove?: () => void
+}
+
+const CollapsibleCodeBlock = ({ onRemove, ...props }: CollapsibleCodeBlockProps) => {
+  const [isExpanded, setIsExpanded] = useState(false)
+
+  const codeString = (props.value || props.children) as string
+  const firstLine = isExpanded
+    ? codeString
+    : codeString?.substring(0, codeString.indexOf('\n')) || codeString
+
+  return (
+    <div className="relative">
+      <div
+        className={cn(
+          'flex items-center gap-1 p-1 bg-surface-100 border border-default w-full overflow-hidden',
+          'rounded-md'
+        )}
+      >
+        <Button
+          type="text"
+          size="tiny"
+          className="w-6 h-6"
+          onClick={() => setIsExpanded(!isExpanded)}
+          icon={isExpanded ? <ChevronUp size={14} /> : <ChevronDown size={14} />}
+        />
+        <div className="flex-1 shrink-1 overflow-auto [&::-webkit-scrollbar]:hidden [-ms-overflow-style:none] [scrollbar-width:none]">
+          <CodeBlock
+            {...props}
+            value={firstLine}
+            hideCopy
+            className={cn(
+              'block !bg-transparent max-h-32 max-w-full !py-0 !px-0 !border-t-0 prose dark:prose-dark border-0 text-foreground !rounded-none w-full text-wrap whitespace-pre-wrap',
+              // change the look of the code block. The flex hack is so that the code is wrapping since
+              // every word is a separate span
+              '[&>code]:m-0 [&>code>span]:flex border-t-0 [&>code>span]:flex-wrap [&>code]:block [&>code>span]:text-foreground text-wrap whitespace-pre-wrap',
+              props.className
+            )}
+          />
+        </div>
+
+        {onRemove && (
+          <Button
+            type="text"
+            size="tiny"
+            className="shrink-0 w-6 h-6"
+            onClick={onRemove}
+            icon={<X size={14} />}
+          />
+        )}
+      </div>
+    </div>
+  )
+}
+
+export default CollapsibleCodeBlock

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

@@ -9,18 +9,18 @@ import { SqlSnippet } from './SqlSnippet'
 import { DiffType } from 'components/interfaces/SQLEditor/SQLEditor.types'
 import { useProfile } from 'lib/profile'
 import { MessagePre } from './MessagePre'
+import { motion } from 'framer-motion'
+import CollapsibleCodeBlock from './CollapsibleCodeBlock'
 
 interface MessageProps {
   name?: string
   role: 'function' | 'system' | 'user' | 'assistant' | 'data' | 'tool'
   content?: string
   createdAt?: number
-  isDebug?: boolean
   isSelected?: boolean
   isLoading?: boolean
+  readOnly?: boolean
   action?: React.ReactNode
-  context?: { entity: string; schemas: string[]; tables: string[] }
-  onDiff?: (type: DiffType, s: string) => void
   variant?: 'default' | 'warning'
 }
 
@@ -29,43 +29,35 @@ export const Message = function Message({
   role,
   content,
   createdAt,
-  isDebug,
   isLoading,
+  readOnly,
   isSelected = false,
-  context,
   children,
   action = null,
   variant = 'default',
-  onDiff = noop,
 }: PropsWithChildren<MessageProps>) {
   const { profile } = useProfile()
   const isUser = role === 'user'
 
-  const formattedContext =
-    context !== undefined
-      ? Object.entries(context)
-          .filter(([_, value]) => value.length > 0)
-          .map(([key, value]) => {
-            return `${key.charAt(0).toUpperCase() + key.slice(1)}: ${Array.isArray(value) ? value.join(', ') : value}`
-          })
-          .join(' • ')
-      : undefined
-
   if (!content) return null
 
   return (
-    <div className="w-full flex flex-col">
+    <motion.div
+      layout="position"
+      initial={{ y: 5, opacity: 0, scale: 0.99 }}
+      animate={{ y: 0, opacity: 1, scale: 1 }}
+      className="w-full flex flex-col"
+    >
       {children}
       <div
         className={cn(
-          'text-foreground-light text-sm mb max-w-full mb-3',
+          'text-foreground-light text-sm mb max-w-full mb-6',
           variant === 'warning' && 'bg-warning-200',
-          isUser ? 'px-4 py-2 rounded-lg bg-background-muted w-fit self-end' : 'mb-6'
+          isUser ? 'px-5 py-3 rounded-lg bg-background-muted w-fit self-end' : 'mb-6'
         )}
       >
         {variant === 'warning' && <WarningIcon className="w-6 h-6" />}
 
-        {isDebug && <Badge variant="warning">Debug request</Badge>}
         {action}
 
         <ReactMarkdown
@@ -74,7 +66,21 @@ export const Message = function Message({
           components={{
             ...markdownComponents,
             pre: (props: any) => {
-              return <SqlSnippet isLoading={isLoading} sql={props.children[0].props.children} />
+              return readOnly ? (
+                <div className="mb-1 -mt-2">
+                  <CollapsibleCodeBlock
+                    value={props.children[0].props.children[0]}
+                    language="sql"
+                    hideLineNumbers
+                  />
+                </div>
+              ) : (
+                <SqlSnippet
+                  readOnly={readOnly}
+                  isLoading={isLoading}
+                  sql={props.children[0].props.children}
+                />
+              )
             },
             ol: (props: any) => {
               return <ol className="flex flex-col gap-y-4">{props.children}</ol>
@@ -92,10 +98,7 @@ export const Message = function Message({
         >
           {content}
         </ReactMarkdown>
-        {/* {role === 'user' && context !== undefined && (
-        <span className="text-xs text-foreground-lighter">{formattedContext}</span>
-      )} */}
       </div>
-    </div>
+    </motion.div>
   )
 }

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

@@ -1,5 +1,5 @@
 import { useState, useEffect, useCallback } from 'react'
-import { Code, Play, Edit, DatabaseIcon } from 'lucide-react'
+import { Code, Play, Edit, DatabaseIcon, FileWarning } from 'lucide-react'
 import { Button, CodeBlock, cn } from 'ui'
 import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
 import { useSelectedProject } from 'hooks/misc/useSelectedProject'
@@ -14,12 +14,14 @@ import useNewQuery from 'components/interfaces/SQLEditor/hooks'
 interface SqlSnippetWrapperProps {
   sql: string
   isLoading?: boolean
+  readOnly?: boolean
 }
 
 interface ParsedSqlProps {
   sql: string
   props: any
   title: string
+  readOnly?: boolean
 }
 
 const isReadOnlySelect = (query: string): boolean => {
@@ -65,7 +67,7 @@ const isReadOnlySelect = (query: string): boolean => {
   })
 }
 
-const SqlSnippetWrapper = ({ sql, isLoading }: SqlSnippetWrapperProps) => {
+const SqlSnippetWrapper = ({ sql, isLoading, readOnly = false }: SqlSnippetWrapperProps) => {
   let formatted = (sql || [''])[0]
   const propsMatch = formatted.match(/--\s*props:\s*(\{[^}]+\})/)
   const props = propsMatch ? JSON.parse(propsMatch[1]) : {}
@@ -96,7 +98,7 @@ const SqlSnippetWrapper = ({ sql, isLoading }: SqlSnippetWrapperProps) => {
     <div
       draggable
       onDragStart={handleDragStart}
-      className="-mx-5 my-3 bg-background-muted border-t border-b overflow-hidden"
+      className="-mx-5 my-3 mt-2 border-b overflow-hidden"
     >
       <SqlCard
         sql={formatted}
@@ -104,17 +106,26 @@ const SqlSnippetWrapper = ({ sql, isLoading }: SqlSnippetWrapperProps) => {
         xAxis={props.xAxis}
         yAxis={props.yAxis}
         title={title}
+        readOnly={readOnly}
       />
     </div>
   )
 }
 
-export const SqlCard = ({ sql, isChart, xAxis, yAxis, title }: ParsedSqlProps) => {
-  const [showCode, setShowCode] = useState(!isReadOnlySelect(sql))
+export const SqlCard = ({
+  sql,
+  isChart,
+  xAxis,
+  yAxis,
+  title,
+  readOnly = false,
+}: ParsedSqlProps) => {
+  const [showCode, setShowCode] = useState(readOnly || !isReadOnlySelect(sql))
   const [showResults, setShowResults] = useState(false)
   const [results, setResults] = useState<any[]>()
   const project = useSelectedProject()
   const [error, setError] = useState<QueryResponseError>()
+  const [showWarning, setShowWarning] = useState(false)
   const { newQuery } = useNewQuery()
 
   const { mutate: executeSql, isLoading: isExecuting } = useExecuteSqlMutation({
@@ -122,17 +133,25 @@ export const SqlCard = ({ sql, isChart, xAxis, yAxis, title }: ParsedSqlProps) =
       console.log('SQL executed successfully:', res)
       setShowResults(true)
       setResults(res.result)
+      setShowWarning(false)
     },
     onError: (error) => {
       setError(error)
       setResults([])
+      setShowWarning(false)
     },
   })
 
   const handleExecute = useCallback(() => {
-    if (!project?.ref || !sql) return
-    setError(undefined)
+    if (!project?.ref || !sql || readOnly) return
 
+    if (!isReadOnlySelect(sql)) {
+      setShowCode(true)
+      setShowWarning(true)
+      return
+    }
+
+    setError(undefined)
     executeSql({
       sql: suffixWithLimit(sql, 100),
       projectRef: project?.ref,
@@ -142,55 +161,97 @@ export const SqlCard = ({ sql, isChart, xAxis, yAxis, title }: ParsedSqlProps) =
         setResults([])
       },
     })
-  }, [project?.ref, project?.connectionString, sql, executeSql])
+  }, [project?.ref, project?.connectionString, sql, executeSql, readOnly])
 
   useEffect(() => {
-    if (isReadOnlySelect(sql) && !results) {
+    if (isReadOnlySelect(sql) && !results && !readOnly) {
       handleExecute()
     }
-  }, [sql, handleExecute])
+  }, [sql, handleExecute, readOnly])
 
   const handleEdit = () => {
-    newQuery(sql, title)
+    if (!readOnly) {
+      newQuery(sql, title)
+    }
   }
 
   const [errorHeader, ...errorContent] =
     (error?.formattedError?.split('\n') ?? [])?.filter((x: string) => x.length > 0) ?? []
 
-  console.log('results:', results)
-
   return (
     <div className="overflow-hidden">
       <div className="flex items-center px-5 py-2 gap-2">
-        <DatabaseIcon size={16} strokeWidth={1.5} />
-        <h3 className="text-sm font-medium flex-1">{title}</h3>
+        {showWarning ? (
+          <div className="py-2">
+            <FileWarning strokeWidth={1.5} size={20} className="text-warning-600 mb-3" />
+            <h3 className="text-sm font-medium flex-1">
+              This query contains write operations. Are you sure you want to execute it?
+            </h3>
+            <div className="flex justify-stretch mt-2 gap-2">
+              <Button
+                type="outline"
+                size="tiny"
+                className="w-full flex-1"
+                onClick={() => setShowWarning(false)}
+              >
+                Cancel
+              </Button>
+              <Button
+                type="outline"
+                size="tiny"
+                className="w-full flex-1"
+                onClick={() => {
+                  setShowWarning(false)
+                  executeSql({
+                    sql: suffixWithLimit(sql, 100),
+                    projectRef: project?.ref,
+                    connectionString: project?.connectionString,
+                    handleError: (error) => {
+                      setError(error)
+                      setResults([])
+                    },
+                  })
+                }}
+              >
+                Run
+              </Button>
+            </div>
+          </div>
+        ) : (
+          <>
+            <DatabaseIcon size={16} strokeWidth={1.5} />
+            <h3 className="text-sm font-medium flex-1">{title}</h3>
 
-        <div className="flex">
-          <Button
-            type="text"
-            size="tiny"
-            className="w-7 h-7"
-            icon={<Code size={14} />}
-            onClick={() => setShowCode(!showCode)}
-          ></Button>
+            {!readOnly && (
+              <div className="flex">
+                <Button
+                  type="text"
+                  size="tiny"
+                  className="w-7 h-7"
+                  icon={<Code size={14} />}
+                  onClick={() => setShowCode(!showCode)}
+                ></Button>
 
-          <Button
-            type="text"
-            size="tiny"
-            className="w-7 h-7"
-            icon={<Edit size={14} />}
-            onClick={handleEdit}
-          ></Button>
+                <Button
+                  type="text"
+                  size="tiny"
+                  className="w-7 h-7"
+                  icon={<Edit size={14} />}
+                  onClick={handleEdit}
+                ></Button>
 
-          <Button
-            type="text"
-            size="tiny"
-            className="w-7 h-7"
-            icon={<Play size={14} />}
-            loading={isExecuting}
-            onClick={handleExecute}
-          ></Button>
-        </div>
+                <Button
+                  type="text"
+                  size="tiny"
+                  className="w-7 h-7"
+                  icon={<Play size={14} />}
+                  loading={isExecuting}
+                  onClick={handleExecute}
+                ></Button>
+              </div>
+            )}
+          </>
+        )}
       </div>
 
       {showCode && (
@@ -199,8 +260,6 @@ export const SqlCard = ({ sql, isChart, xAxis, yAxis, title }: ParsedSqlProps) =
           language="sql"
           className={cn(
             'max-h-96 block !bg-transparent !py-3 !px-3.5 prose dark:prose-dark border-0 border-t text-foreground !rounded-none w-full',
-            // change the look of the code block. The flex hack is so that the code is wrapping since
-            // every word is a separate span
             '[&>code]:m-0 [&>code>span]:flex [&>code>span]:flex-wrap [&>code]:block [&>code>span]:text-foreground'
           )}
           hideLineNumbers

apps/studio/pages/api/ai/sql/generate-v3.ts

@@ -59,7 +59,11 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
 
       # You write row level security policies.
 
-      Your purpose is to generate a policy with the constraints given by the user. First retrieve the schema or schemas that are most relevant. Then retrieve RLS policy examples.
+      Your purpose is to generate a policy with the constraints given by the user.
+      - First, use getSchema to retrieve more information about a schema or schemas that will contain policies, usually the public schema.
+      - Then retrieve existing RLS policies and guidelines on how to write policies using the getRlsKnowledge tool .
+      - Then write new policies or update existing policies based on the prompt
+      - When asked to suggest policies, either alter existing policies or add new ones to the public schema.
 
       # You write database functions
       Your purpose is to generate a database function with the constraints given by the user. The output may also include a database trigger
@@ -72,7 +76,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
       Follow these instructions:
       - First look at the list of provided schemas and if needed, get more information about a schema. You will almost always need to retrieve information about the public schema before answering a question.
 
-      Here are the existing database schemas: ${schemas}
+      Here are the existing database schema names you can retrieve: ${schemas}
       `,
     messages,
     tools: getTools(projectRef, connectionString, authorization),

apps/studio/pages/api/ai/sql/tools.ts

@@ -4,12 +4,12 @@ import { executeSql } from 'data/sql/execute-sql-query'
 import { getEntityDefinitionsQuery } from 'data/database/entity-definitions-query'
 import { getDatabasePolicies } from 'data/database-policies/database-policies-query'
 import { stripIndent } from 'common-tags'
-type DatabasePoliciesData = components['schemas']['PostgresPolicy']
+import { DatabasePoliciesData } from 'data/database-policies/database-policies-query'
 
 export const getTools = (projectRef: string, connectionString: string, authorization?: string) => {
   return {
     getSchema: tool({
-      description: 'Get the schema definition for a given schema',
+      description: 'Get more information about one or more schemas',
       parameters: z.object({
         schemas: z.array(z.string()).describe('The schema names to get the definitions for'),
       }),
@@ -38,8 +38,9 @@ export const getTools = (projectRef: string, connectionString: string, authoriza
         }
       },
     }),
-    getRlsExamples: tool({
-      description: 'Get examples and instructions on how to write RLS policies',
+    getRlsKnowledge: tool({
+      description:
+        'Get existing policies and examples and instructions on how to write RLS policies',
       parameters: z.object({
         schemas: z.array(z.string()).describe('The schema names to get the policies for'),
       }),
@@ -55,7 +56,6 @@ export const getTools = (projectRef: string, connectionString: string, authoriza
             ...(authorization && { Authorization: authorization }),
           }
         )
-        console.log('rls:', data)
         const formattedPolicies = data
           .map(
             (policy: DatabasePoliciesData) => `
@@ -71,8 +71,7 @@ export const getTools = (projectRef: string, connectionString: string, authoriza
 
         return stripIndent`
           You're a Supabase Postgres expert in writing row level security policies. Your purpose is to
-          generate a policy with the constraints given by the user. You will be provided a schema
-          on which the policy should be applied.
+          generate a policy with the constraints given by the user. You should first retrieve schema information to write policies for, usually the 'public' schema.
 
           The output should use the following instructions:
           - The generated SQL must be valid SQL.

apps/studio/pages/project/[ref]/auth/policies.tsx

@@ -21,6 +21,7 @@ import { EXCLUDED_SCHEMAS } from 'lib/constants/schemas'
 import type { NextPageWithLayout } from 'types'
 import { Button, Input } from 'ui'
 import { DocsButton } from 'components/ui/DocsButton'
+import { useAppStateSnapshot } from 'state/app-state'
 
 /**
  * Filter tables by table name and policy name
@@ -64,6 +65,7 @@ const AuthPoliciesPage: NextPageWithLayout = () => {
   }>()
   const { schema = 'public', search: searchString = '' } = params
   const { project } = useProjectContext()
+  const { setAiAssistantPanel, aiAssistantPanel } = useAppStateSnapshot()
 
   const [selectedTable, setSelectedTable] = useState<string>()
   const [showPolicyAiEditor, setShowPolicyAiEditor] = useState(false)
@@ -148,8 +150,10 @@ const AuthPoliciesPage: NextPageWithLayout = () => {
           hasTables={tables.length > 0}
           isLocked={isLocked}
           onSelectCreatePolicy={(table: string) => {
-            setShowPolicyAiEditor(true)
-            setSelectedTable(table)
+            setAiAssistantPanel({
+              open: true,
+              initialInput: `Help me create an RLS policy for the table ${table} on the ${schema} schema`,
+            })
           }}
           onSelectEditPolicy={(policy) => {
             setSelectedPolicyToEdit(policy)

apps/studio/state/app-state.ts

@@ -3,6 +3,7 @@ import { proxy, snapshot, useSnapshot } from 'valtio'
 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
@@ -21,6 +22,9 @@ export type CommonDatabaseEntity = {
 
 type AiAssistantPanelType = {
   open: boolean
+  initialMessages?: MessageType[]
+  initialInput?: string
+  sqlSnippets?: string[]
   editor?: SupportedAssistantEntities | null
   // Raw string content for the monaco editor, currently used to retain where the user left off when toggling off the panel
   content?: string
@@ -109,6 +113,9 @@ export const appState = proxy({
 
   aiAssistantPanel: {
     open: false,
+    initialMessages: [],
+    sqlSnippets: [],
+    initialInput: '',
     editor: null,
     content: '',
     entity: undefined,