Feat/assistant additions (supabase#30647)

* show thinking if last message is empty

* sql toggle copy and border fix

* add table and schema context

* add suggestions for policy

* Add snippet content to ai assistant if on sql editor and assistant is open

---------

Co-authored-by: Joshen Lim <[email protected]>

Author: SaxonF

apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyRow.tsx

@@ -118,6 +118,14 @@ const PolicyRow = ({
                       open: true,
                       sqlSnippets: [sql],
                       initialInput: `Update the policy with name "${policy.name}" in the ${policy.schema} schema on the ${policy.table} table. It should...`,
+                      suggestions: {
+                        title: `I can help you make a change to the policy "${policy.name}" in the ${policy.schema} schema on the ${policy.table} table, here are a few example prompts to get you started:`,
+                        prompts: [
+                          'Tell me how I can improve this policy...',
+                          'Duplicate this policy for another table...',
+                          'Add extra conditions to this policy...',
+                        ],
+                      },
                     })
                   }}
                 >

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

@@ -40,6 +40,10 @@ import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
 import AIOnboarding from './AIOnboarding'
 import CollapsibleCodeBlock from './CollapsibleCodeBlock'
 import { Message } from './Message'
+import { useParams } from 'common/hooks'
+import { useSearchParamsShallow } from 'common/hooks'
+import { useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2'
+import { useRouter } from 'next/router'
 
 const MemoizedMessage = memo(
   ({ message, isLoading }: { message: MessageType; isLoading: boolean }) => {
@@ -71,14 +75,18 @@ export const AIAssistant = ({
   className,
   onResetConversation,
 }: AIAssistantProps) => {
+  const router = useRouter()
   const project = useSelectedProject()
   const isOptedInToAI = useOrgOptedIntoAi()
   const selectedOrganization = useSelectedOrganization()
+  const { id: entityId } = useParams()
+  const searchParams = useSearchParamsShallow()
   const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM
 
   const disablePrompts = useFlag('disableAssistantPrompts')
+  const { snippets } = useSqlEditorV2StateSnapshot()
   const { aiAssistantPanel, setAiAssistantPanel } = useAppStateSnapshot()
-  const { initialInput, sqlSnippets, suggestions } = aiAssistantPanel
+  const { open, initialInput, sqlSnippets, suggestions } = aiAssistantPanel
 
   const inputRef = useRef<HTMLTextAreaElement>(null)
   const bottomRef = useRef<HTMLDivElement>(null)
@@ -93,6 +101,10 @@ export const AIAssistant = ({
   const { data: check } = useCheckOpenAIKeyQuery()
   const isApiKeySet = IS_PLATFORM || !!check?.hasKey
 
+  const isInSQLEditor = router.pathname.includes('/sql/[id]')
+  const snippet = snippets[entityId ?? '']
+  const snippetContent = snippet?.snippet?.content?.sql
+
   const { data: subscription } = useOrgSubscriptionQuery({ orgSlug: selectedOrganization?.slug })
   const hasHipaaAddon = subscriptionHasHipaaAddon(subscription)
 
@@ -102,6 +114,9 @@ export const AIAssistant = ({
     schema: 'public',
   })
 
+  const currentTable = tables?.find((t) => t.id.toString() === entityId)
+  const currentSchema = searchParams?.get('schema') ?? 'public'
+
   const { mutate: sendEvent } = useSendEventMutation()
   const sendTelemetryEvent = (value: string) => {
     sendEvent({
@@ -126,6 +141,8 @@ export const AIAssistant = ({
       includeSchemaMetadata,
       projectRef: project?.ref,
       connectionString: project?.connectionString,
+      schema: currentSchema,
+      table: currentTable?.name,
     },
   })
 
@@ -264,6 +281,20 @@ export const AIAssistant = ({
     }
   }, [messages, isChatLoading, setAiAssistantPanel])
 
+  // Remove suggestions if sqlSnippets were removed
+  useEffect(() => {
+    if (!sqlSnippets || sqlSnippets.length === 0) {
+      setAiAssistantPanel({ suggestions: undefined })
+    }
+  }, [sqlSnippets, suggestions, setAiAssistantPanel])
+
+  useEffect(() => {
+    if (open && isInSQLEditor && !!snippetContent) {
+      setAiAssistantPanel({ sqlSnippets: [snippetContent] })
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [open, isInSQLEditor, snippetContent])
+
   if (isLoadingTables) {
     return (
       <div className="h-full w-full flex justify-center items-center">
@@ -329,7 +360,7 @@ export const AIAssistant = ({
                 })}
               </div>
               {renderedMessages}
-              {last(messages)?.role === 'user' && (
+              {(last(messages)?.role === 'user' || last(messages)?.content?.length === 0) && (
                 <motion.div className="text-foreground-lighter text-sm flex gap-1.5 items-center">
                   <span>Thinking</span>
                   <div className="flex gap-1">
@@ -363,7 +394,7 @@ export const AIAssistant = ({
               </h3>
               {suggestions.title && <p>{suggestions.title}</p>}
               <div className="-mx-3 mt-4 mb-12">
-                {suggestions?.prompts.map((prompt) => (
+                {suggestions?.prompts?.map((prompt) => (
                   <Button
                     size="small"
                     icon={<FileText strokeWidth={1.5} size={16} />}

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

@@ -222,7 +222,9 @@ export const SqlCard = ({
                   className="w-7 h-7"
                   icon={<Code size={14} />}
                   onClick={() => setShowCode(!showCode)}
-                  tooltip={{ content: { side: 'bottom', text: 'Show query' } }}
+                  tooltip={{
+                    content: { side: 'bottom', text: showCode ? 'Hide query' : 'Show query' },
+                  }}
                 />
 
                 {!isInSQLEditor || isInNewSnippet ? (
@@ -314,7 +316,7 @@ export const SqlCard = ({
           value={sql}
           language="sql"
           className={cn(
-            'max-w-full 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',
+            'max-h-96 max-w-none block !bg-transparent !py-3 !px-3.5 prose dark:prose-dark border-0 border-t text-foreground !rounded-none w-full',
             '[&>code]:m-0 [&>code>span]:flex [&>code>span]:flex-wrap [&>code]:block [&>code>span]:text-foreground'
           )}
         />

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

@@ -39,7 +39,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
 }
 
 async function handlePost(req: NextApiRequest, res: NextApiResponse) {
-  const { messages, projectRef, connectionString, includeSchemaMetadata } = req.body
+  const { messages, projectRef, connectionString, includeSchemaMetadata, schema, table } = req.body
 
   if (!projectRef) {
     return res.status(400).json({
@@ -109,6 +109,9 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
       - 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. If the question is about users, also retrieve the auth schema.
 
       Here are the existing database schema names you can retrieve: ${schemas}
+
+      ${schema !== undefined ? `The user is currently looking at the ${schema} schema.` : ''}
+      ${table !== undefined ? `The user is currently looking at the ${table} table.` : ''}
       `,
     messages,
     tools: getTools({ projectRef, connectionString, authorization, includeSchemaMetadata }),

apps/studio/state/app-state.ts

@@ -22,7 +22,7 @@ export type CommonDatabaseEntity = {
 
 export type SuggestionsType = {
   title: string
-  prompts: string[]
+  prompts?: string[]
 }
 
 type AiAssistantPanelType = {