feat(cmdk): generate titles for sql ai code snippets

package-lock.json

@@ -39064,6 +39064,7 @@
         "clipboard": "^2.0.8",
         "clsx": "^1.2.1",
         "common": "*",
+        "common-tags": "^1.8.2",
         "config": "*",
         "configcat-js": "^7.0.0",
         "dayjs": "^1.11.0",
@@ -39138,6 +39139,7 @@
         "@testing-library/react-hooks": "^8.0.1",
         "@testing-library/user-event": "^13.5.0",
         "@types/blueimp-md5": "^2.18.0",
+        "@types/common-tags": "^1.8.1",
         "@types/file-saver": "^2.0.2",
         "@types/json-logic-js": "^1.2.1",
         "@types/lodash": "^4.14.172",
@@ -66062,6 +66064,7 @@
         "@testing-library/react-hooks": "^8.0.1",
         "@testing-library/user-event": "^13.5.0",
         "@types/blueimp-md5": "^2.18.0",
+        "@types/common-tags": "^1.8.1",
         "@types/file-saver": "^2.0.2",
         "@types/json-logic-js": "^1.2.1",
         "@types/lodash": "^4.14.172",
@@ -66092,6 +66095,7 @@
         "clipboard": "^2.0.8",
         "clsx": "^1.2.1",
         "common": "*",
+        "common-tags": "^1.8.2",
         "config": "*",
         "configcat-js": "^7.0.0",
         "dayjs": "^1.11.0",

packages/ui/src/components/Command/AiCommand.tsx

@@ -272,6 +272,70 @@ export function useAiChat({
   }
 }
 
+/**
+ * Perform a one-off query to AI based on a snapshot of messages
+ */
+export function queryAi(messages: Message[], timeout = 0) {
+  return new Promise<string>((resolve, reject) => {
+    const eventSource = new SSE(`${edgeFunctionUrl}/ai-docs`, {
+      headers: {
+        apikey: process.env.NEXT_PUBLIC_SUPABASE_ANON_KEY ?? '',
+        Authorization: `Bearer ${process.env.NEXT_PUBLIC_SUPABASE_ANON_KEY}`,
+        'Content-Type': 'application/json',
+      },
+      payload: JSON.stringify({
+        messages: messages.map(({ role, content }) => ({ role, content })),
+      }),
+    })
+
+    let timeoutId: number | undefined
+
+    function handleError<T>(err: T) {
+      if (timeoutId) {
+        clearTimeout(timeoutId)
+      }
+      console.error(err)
+      reject(err)
+    }
+
+    if (timeout > 0) {
+      timeoutId = window.setTimeout(() => {
+        handleError(new Error('AI query timed out'))
+      }, timeout)
+    }
+
+    let answer = ''
+
+    eventSource.addEventListener('error', handleError)
+    eventSource.addEventListener('message', (e) => {
+      try {
+        if (e.data === '[DONE]') {
+          if (timeoutId) {
+            clearTimeout(timeoutId)
+          }
+          resolve(answer)
+          return
+        }
+
+        const completionResponse: CreateChatCompletionResponse = JSON.parse(e.data)
+        const [
+          {
+            delta: { content },
+          },
+        ] = completionResponse.choices as CreateChatCompletionResponseChoicesInnerDelta[]
+
+        if (content) {
+          answer += content
+        }
+      } catch (err) {
+        handleError(err)
+      }
+    })
+
+    eventSource.stream()
+  })
+}
+
 const AiCommand = () => {
   const { isLoading, setIsLoading, search, setSearch, MarkdownHandler } = useCommandMenu()
 

packages/ui/src/components/Command/CommandMenuProvider.tsx

@@ -25,7 +25,7 @@ export interface CommandMenuContextValue {
   MarkdownHandler: (props: ReactMarkdownOptions) => JSX.Element // to do: remove this. although it breaks our jest tests
 
   // Optional callback to save a generated SQL output
-  onSaveGeneratedSQL?: (answer: string, resolve: any) => void
+  saveGeneratedSQL?: (answer: string, title: string) => Promise<void>
 }
 export const CommandMenuContext = createContext<CommandMenuContextValue | undefined>(undefined)
 export const useCommandMenu = () => {
@@ -49,7 +49,7 @@ export interface CommandMenuProviderProps {
   // to do: remove this prop
   // this is a temporary hack as ReactMarkdown fails our jest tests if we import the package within this UI package
   MarkdownHandler: (props: ReactMarkdownOptions) => JSX.Element
-  onSaveGeneratedSQL?: (answer: string, resolve: any) => void
+  saveGeneratedSQL?: (answer: string, title: string) => Promise<void>
 }
 
 const CommandMenuProvider = ({
@@ -58,7 +58,7 @@ const CommandMenuProvider = ({
   projectRef,
   apiKeys,
   MarkdownHandler,
-  onSaveGeneratedSQL,
+  saveGeneratedSQL,
 }: PropsWithChildren<CommandMenuProviderProps>) => {
   const [isOpen, setIsOpen] = useState(false)
   const [isLoading, setIsLoading] = useState(false)
@@ -88,7 +88,7 @@ const CommandMenuProvider = ({
         site,
         project,
         MarkdownHandler,
-        onSaveGeneratedSQL,
+        saveGeneratedSQL,
       }}
     >
       {children}

packages/ui/src/components/Command/GenerateSQL.tsx

@@ -10,34 +10,77 @@ import {
   IconSave,
   IconUser,
   Input,
+  Message,
   MessageRole,
   MessageStatus,
+  queryAi,
   useAiChat,
 } from 'ui'
 
+import { stripIndent } from 'common-tags'
 import CopyToClipboard from 'react-copy-to-clipboard'
 import { cn } from './../../utils/cn'
 import { SAMPLE_QUERIES } from './Command.constants'
 import { AiIcon, AiIconChat } from './Command.icons'
 import { CommandItem } from './Command.utils'
 import { useCommandMenu } from './CommandMenuProvider'
-import { stripIndent } from 'common-tags'
 
-const SQLOutputActions = ({ answer }: { answer: string }) => {
+/**
+ * Formats a string for use as a title.
+ *
+ * Removes punctuation and capitalizes each word
+ */
+function formatTitle(value: string) {
+  let words = value.replace(/\.$/, '').replace(/['"]/g, '').split(' ')
+  words = words.map((word) => {
+    // Don't capitalize code
+    if (/[._\(\)]+/.test(word)) {
+      return word
+    }
+    return word.charAt(0).toUpperCase() + word.slice(1)
+  })
+  return words.join(' ')
+}
+
+const SQLOutputActions = ({ answer, messages }: { answer: string; messages: Message[] }) => {
   const [showCopied, setShowCopied] = useState(false)
   const [isSaving, setIsSaving] = useState(false)
   const [isSaved, setIsSaved] = useState(false)
 
-  const { project, onSaveGeneratedSQL } = useCommandMenu()
-
-  const applyCallback = () =>
-    onSaveGeneratedSQL !== undefined
-      ? new Promise((resolve) => onSaveGeneratedSQL(answer, resolve))
-      : {}
+  const { project, saveGeneratedSQL } = useCommandMenu()
 
   const onSelectSaveSnippet = async () => {
     setIsSaving(true)
-    await applyCallback()
+    let suggestedTitle
+    try {
+      suggestedTitle = await queryAi(
+        [
+          ...messages,
+          {
+            role: MessageRole.User,
+            content: stripIndent`
+            Generate a title for the above SQL snippet following all of these rules:
+            - The title is only for the last SQL snippet
+            - Focus on the main purposes of this snippet
+            - Use as few words as possible
+            - Title should be nouns, not verbs
+            - Do not include word articles (eg. a, the, for, of)
+            - Do not use words like "SQL" or "snippet"
+            - Do not output markdown, quotes, etc
+            - Do not be too verbose
+            `,
+            status: MessageStatus.Complete,
+          },
+        ],
+        10000
+      )
+    } catch (err) {
+      suggestedTitle = ''
+    }
+
+    const formattedTitle = formatTitle(suggestedTitle)
+
+    await saveGeneratedSQL?.(answer, formattedTitle)
     setIsSaved(true)
     setIsSaving(false)
   }
@@ -71,7 +114,7 @@ const SQLOutputActions = ({ answer }: { answer: string }) => {
           {showCopied ? 'Copied' : 'Copy SQL'}
         </Button>
       </CopyToClipboard>
-      {project?.ref !== undefined && onSaveGeneratedSQL !== undefined && (
+      {project?.ref !== undefined && saveGeneratedSQL !== undefined && (
         <Button
           type="default"
           loading={isSaving}
@@ -147,7 +190,7 @@ const GenerateSQL = () => {
   return (
     <div onClick={(e) => e.stopPropagation()}>
       <div className={cn('relative mb-[62px] py-4 max-h-[720px] overflow-auto')}>
-        {messages.map((message) => {
+        {messages.map((message, i) => {
           switch (message.role) {
             case MessageRole.User:
               return (
@@ -206,7 +249,9 @@ const GenerateSQL = () => {
                           >
                             {answer}
                           </CodeBlock>
-                          {!isResponding && <SQLOutputActions answer={answer} />}
+                          {!isResponding && (
+                            <SQLOutputActions answer={answer} messages={messages.slice(0, i + 1)} />
+                          )}
                         </div>
                       )}
                     </>

studio/components/interfaces/App/CommandMenuWrapper.tsx

@@ -13,6 +13,7 @@ import { useSqlEditorStateSnapshot } from 'state/sql-editor'
 import { useProfileQuery } from 'data/profile/profile-query'
 import { SqlSnippet } from 'data/content/sql-snippets-query'
 import { useProjectApiQuery } from 'data/config/project-api-query'
+import { codeBlock } from 'common-tags'
 
 const CommandMenuWrapper = observer(({ children }: PropsWithChildren<{}>) => {
   const { ref } = useParams()
@@ -31,30 +32,30 @@ const CommandMenuWrapper = observer(({ children }: PropsWithChildren<{}>) => {
     service: settings?.autoApiService?.serviceApiKey ?? undefined,
   }
 
-  const onSaveGeneratedSQL = async (answer: string, resolve: any) => {
+  const onSaveGeneratedSQL = async (answer: string, title: string) => {
     if (!ref) return console.error('Project ref is required')
     if (!canCreateSQLSnippet) {
       ui.setNotification({
         category: 'info',
         message: 'Unable to save query as you do not have sufficient permissions for this project',
       })
-      return resolve()
+      return
     }
 
     // Remove markdown syntax from returned answer
     answer = answer.replace(/`/g, '').replace(/sql\n/g, '').trim()
 
-    const formattedSql = `
--- Note: This query was generated via Supabase AI, please verify the correctness of the
--- SQL snippet before running it against your database as we are not able to guarantee it
--- will do exactly what you requested the AI.
+    const formattedSql = codeBlock`
+      -- Note: This query was generated via Supabase AI, please verify the correctness of the
+      -- SQL snippet before running it against your database as we are not able to guarantee it
+      -- will do exactly what you requested the AI.
 
-${answer}
-`.trim()
+      ${answer}
+    `
 
     try {
       const snippet = createSqlSnippetSkeleton({
-        name: 'Generated query',
+        name: title || 'Generated query',
         owner_id: profile?.id,
         sql: formattedSql,
       })
@@ -69,8 +70,6 @@ ${answer}
         category: 'error',
         message: `Failed to create new query: ${error.message}`,
       })
-    } finally {
-      resolve()
     }
   }
 
@@ -80,7 +79,7 @@ ${answer}
       projectRef={ref}
       apiKeys={apiKeys}
       MarkdownHandler={(props) => <ReactMarkdown remarkPlugins={[remarkGfm]} {...props} />}
-      onSaveGeneratedSQL={onSaveGeneratedSQL}
+      saveGeneratedSQL={onSaveGeneratedSQL}
     >
       {children}
     </CommandMenuProvider>

studio/package.json

@@ -43,6 +43,7 @@
     "clipboard": "^2.0.8",
     "clsx": "^1.2.1",
     "common": "*",
+    "common-tags": "^1.8.2",
     "config": "*",
     "configcat-js": "^7.0.0",
     "dayjs": "^1.11.0",
@@ -117,6 +118,7 @@
     "@testing-library/react-hooks": "^8.0.1",
     "@testing-library/user-event": "^13.5.0",
     "@types/blueimp-md5": "^2.18.0",
+    "@types/common-tags": "^1.8.1",
     "@types/file-saver": "^2.0.2",
     "@types/json-logic-js": "^1.2.1",
     "@types/lodash": "^4.14.172",