Implement generate sql with AI in CMDK for studio

packages/ui/src/components/Command/Command.constants.ts

@@ -1,6 +1,7 @@
 export const COMMAND_ROUTES = {
   AI: 'Supabase AI',
   DOCS_SEARCH: 'Docs Search',
+  GENERATE_SQL: 'Generate SQL',
   THEME: 'Theme',
   AI_ASK_ANYTHING: 'Ask anything',
   AI_RLS_POLICY: 'Help me create a RLS policy',

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

@@ -1,12 +1,17 @@
 import React from 'react'
-export const AiIcon = () => (
+
+interface AiIconProps {
+  className?: string
+}
+
+export const AiIcon = ({ className = 'text-brand-900' }: AiIconProps) => (
   <svg
     xmlns="http://www.w3.org/2000/svg"
     fill="none"
     viewBox="0 0 24 24"
     strokeWidth="1.5"
     stroke="currentColor"
-    className="w-6 h-6 text-brand-900"
+    className={`w-6 h-6 ${className}`}
   >
     <path
       strokeLinecap="round"

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

@@ -32,11 +32,13 @@ import CommandMenuShortcuts from './CommandMenuShortcuts'
 import SearchOnlyItem from './SearchOnlyItem'
 import SearchableStudioItems from './SearchableStudioItems'
 import { COMMAND_ROUTES } from './Command.constants'
+import GenerateSQL from './GenerateSQL'
 
 export const CHAT_ROUTES = [
   COMMAND_ROUTES.AI, // this one is temporary
   COMMAND_ROUTES.AI_ASK_ANYTHING,
   COMMAND_ROUTES.AI_RLS_POLICY,
+  COMMAND_ROUTES.GENERATE_SQL,
 ]
 
 const iconPicker: { [key: string]: React.ReactNode } = {
@@ -193,10 +195,22 @@ const CommandMenu = ({ projectRef }: CommandMenuProps) => {
                 </CommandGroup>
               )}
 
-              {site === 'studio' && (
+              {site === 'studio' && [
+                <CommandGroup heading="Experimental">
+                  <CommandItem
+                    forceMount
+                    type="command"
+                    onSelect={() => setPages([...pages, COMMAND_ROUTES.GENERATE_SQL])}
+                  >
+                    <AiIcon className="text-scale-1100" />
+                    <CommandLabel>Generate SQL with Supabase AI</CommandLabel>
+                  </CommandItem>
+                </CommandGroup>,
                 <CommandGroup heading="Navigate">
                   {sharedItems.tools.map((item) => {
-                    const itemUrl = projectRef ? item.url.replace('_', projectRef) : item.url
+                    const itemUrl = (
+                      projectRef ? item.url.replace('_', projectRef) : item.url
+                    ).split('https://app.supabase.com')[1]
 
                     return (
                       <CommandItem key={item.url} type="link" onSelect={() => router.push(itemUrl)}>
@@ -207,8 +221,8 @@ const CommandMenu = ({ projectRef }: CommandMenuProps) => {
                       </CommandItem>
                     )
                   })}
-                </CommandGroup>
-              )}
+                </CommandGroup>,
+              ]}
 
               {/* <DashboardTableEditor /> */}
 
@@ -244,6 +258,7 @@ const CommandMenu = ({ projectRef }: CommandMenuProps) => {
           )}
           {currentPage === COMMAND_ROUTES.AI && <AiCommand />}
           {currentPage === COMMAND_ROUTES.DOCS_SEARCH && <DocsSearch />}
+          {currentPage === COMMAND_ROUTES.GENERATE_SQL && <GenerateSQL />}
           {currentPage === COMMAND_ROUTES.THEME && <ThemeOptions />}
         </CommandList>
       </CommandDialog>

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

@@ -19,6 +19,9 @@ export interface CommandMenuContextValue {
   // 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: any) => React.ReactNode // to do: remove this. although it breaks our jest tests
+
+  // Optional callback to save a generated SQL output
+  onSaveGeneratedSQL?: (answer: string, resolve: any) => void
 }
 export const CommandMenuContext = createContext<CommandMenuContextValue | undefined>(undefined)
 export const useCommandMenu = () => {
@@ -41,6 +44,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: any) => React.ReactNode
+  onSaveGeneratedSQL?: (answer: string, resolve: any) => void
 }
 
 const CommandMenuProvider = ({
@@ -48,6 +52,7 @@ const CommandMenuProvider = ({
   site,
   projectRef,
   MarkdownHandler,
+  onSaveGeneratedSQL,
 }: PropsWithChildren<CommandMenuProviderProps>) => {
   const [isOpen, setIsOpen] = useState(false)
   const [isLoading, setIsLoading] = useState(false)
@@ -77,6 +82,7 @@ const CommandMenuProvider = ({
         currentPage,
         site,
         MarkdownHandler,
+        onSaveGeneratedSQL,
       }}
     >
       {children}

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

@@ -0,0 +1,417 @@
+import { SSE } from 'sse.js'
+import type { CreateCompletionResponse } from 'openai'
+import { useCallback, useEffect, useReducer, useRef, useState } from 'react'
+import {
+  Button,
+  CodeBlock,
+  IconAlertCircle,
+  IconAlertTriangle,
+  IconCheck,
+  IconClipboard,
+  IconCornerDownLeft,
+  IconSave,
+  IconUser,
+  Input,
+} from 'ui'
+
+import { cn } from './../../utils/cn'
+import { AiIcon, AiIconChat } from './Command.icons'
+import { CommandGroup, CommandItem } from './Command.utils'
+import { useCommandMenu } from './CommandMenuProvider'
+import CopyToClipboard from 'react-copy-to-clipboard'
+import { noop } from 'lodash'
+
+const SAMPLE_QUERIES = [
+  'Create a table that stores a list of cities, and insert 10 rows of sample data into it',
+  'Generate tables (with id bigserial & FK relationships) for blog posts and comments',
+  'Create a trigger that updates the updated_at column on the orders table with the current time when the row of the orders table is updated',
+]
+
+function getEdgeFunctionUrl() {
+  const supabaseUrl = process.env.NEXT_PUBLIC_SUPABASE_URL?.replace(/\/$/, '')
+
+  if (!supabaseUrl) return undefined
+
+  // https://github.com/supabase/supabase-js/blob/10d3423506cbd56345f7f6ab2ec2093c8db629d4/src/SupabaseClient.ts#L96
+  const isPlatform = supabaseUrl.match(/(supabase\.co)|(supabase\.in)/)
+
+  if (isPlatform) {
+    const [schemeAndProjectId, domain, tld] = supabaseUrl.split('.')
+    return `${schemeAndProjectId}.functions.${domain}.${tld}`
+  } else {
+    return `${supabaseUrl}/functions/v1`
+  }
+}
+
+function promptDataReducer(
+  state: any[],
+  action: {
+    index?: number
+    answer?: string | undefined
+    status?: string
+    query?: string | undefined
+    type?: 'remove-last-item' | string
+  }
+) {
+  // set a standard state to use later
+  let current = [...state]
+
+  if (action.type) {
+    switch (action.type) {
+      case 'remove-last-item':
+        current.pop()
+        return [...current]
+      default:
+        break
+    }
+  }
+
+  // check that an index is present
+  if (action.index === undefined) return [...state]
+
+  if (!current[action.index]) {
+    current[action.index] = { query: '', answer: '', status: '' }
+  }
+
+  current[action.index].answer = action.answer
+
+  if (action.query) {
+    current[action.index].query = action.query
+  }
+  if (action.status) {
+    current[action.index].status = action.status
+  }
+
+  return [...current]
+}
+
+const SQLOutputActions = ({
+  answer,
+  onSaveGeneratedSQL,
+}: {
+  answer: string
+  onSaveGeneratedSQL?: (answer: string, resolve: any) => void
+}) => {
+  const [showCopied, setShowCopied] = useState(false)
+  const [isSaving, setIsSaving] = useState(false)
+  const [isSaved, setIsSaved] = useState(false)
+
+  const applyCallback = () =>
+    onSaveGeneratedSQL !== undefined
+      ? new Promise((resolve) => onSaveGeneratedSQL(answer, resolve))
+      : {}
+
+  const onSelectSaveSnippet = async () => {
+    setIsSaving(true)
+    await applyCallback()
+    setIsSaved(true)
+    setIsSaving(false)
+  }
+
+  useEffect(() => {
+    if (!showCopied) return
+    const timer = setTimeout(() => setShowCopied(false), 2000)
+    return () => clearTimeout(timer)
+  }, [showCopied])
+
+  useEffect(() => {
+    if (!isSaved) return
+    const timer = setTimeout(() => setIsSaved(false), 2000)
+    return () => clearTimeout(timer)
+  }, [isSaved])
+
+  return (
+    <div className="flex items-center justify-end space-x-2 mr-12">
+      <CopyToClipboard text={answer}>
+        <Button
+          type="default"
+          icon={
+            showCopied ? (
+              <IconCheck size="tiny" className="text-brand-900" strokeWidth={2} />
+            ) : (
+              <IconClipboard size="tiny" />
+            )
+          }
+          onClick={() => setShowCopied(true)}
+        >
+          {showCopied ? 'Copied' : 'Copy SQL'}
+        </Button>
+      </CopyToClipboard>
+      {onSaveGeneratedSQL !== undefined && (
+        <Button
+          type="default"
+          loading={isSaving}
+          disabled={isSaving}
+          icon={
+            isSaved ? (
+              <IconCheck size="tiny" className="text-brand-900" strokeWidth={2} />
+            ) : (
+              <IconSave size="tiny" />
+            )
+          }
+          onClick={() => onSelectSaveSnippet()}
+        >
+          {isSaved ? 'Snippet saved!' : 'Save into new snippet'}
+        </Button>
+      )}
+    </div>
+  )
+}
+
+const GenerateSQL = () => {
+  const [promptIndex, setPromptIndex] = useState(0)
+  const [answer, setAnswer] = useState<string | undefined>('')
+  const [isResponding, setIsResponding] = useState(false)
+  const [hasClippyError, setHasClippyError] = useState(false)
+
+  const eventSourceRef = useRef<SSE>()
+  const [promptData, dispatchPromptData] = useReducer(promptDataReducer, [])
+  const { isLoading, setIsLoading, search, setSearch, MarkdownHandler, onSaveGeneratedSQL } =
+    useCommandMenu()
+
+  const cantHelp = answer?.trim() === "Sorry, I don't know how to help with that."
+
+  const handleConfirm = useCallback(
+    async (query: string) => {
+      const edgeFunctionUrl = getEdgeFunctionUrl()
+
+      if (!edgeFunctionUrl) {
+        return console.error('No edge function url')
+      }
+
+      setAnswer(undefined)
+      setSearch('')
+      dispatchPromptData({ index: promptIndex, answer: undefined, query })
+      setIsResponding(false)
+      setHasClippyError(false)
+      setIsLoading(true)
+
+      // [Joshen] Eventually we need to pass the table data in here as well
+      const queryToSend = `
+Generate a Postgres SQL query based on the following natural language prompt:
+${query}
+
+Postgres SQL query:
+`.trim()
+
+      const eventSource = new SSE(`${edgeFunctionUrl}/clippy-search`, {
+        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({ query: queryToSend, context: promptData }),
+      })
+
+      function handleError<T>(err: T) {
+        setIsLoading(false)
+        setIsResponding(false)
+        setHasClippyError(true)
+        console.error(err)
+      }
+
+      eventSource.addEventListener('error', handleError)
+      eventSource.addEventListener('message', (e: any) => {
+        try {
+          setIsLoading(false)
+          if (e.data === '[DONE]') {
+            setIsResponding(false)
+            setAnswer(undefined)
+            setPromptIndex((x) => {
+              return x + 1
+            })
+            return
+          }
+          setIsResponding(true)
+
+          const completionResponse: CreateCompletionResponse = JSON.parse(e.data)
+          const [{ text: content }] = completionResponse.choices
+          const text = content ?? ''
+
+          setAnswer((answer) => {
+            const currentAnswer = answer ?? ''
+            dispatchPromptData({ index: promptIndex, answer: currentAnswer + text })
+            return (answer ?? '') + text
+          })
+        } catch (err) {
+          handleError(err)
+        }
+      })
+
+      eventSource.stream()
+      eventSourceRef.current = eventSource
+      setIsLoading(true)
+    },
+    [promptIndex, promptData]
+  )
+
+  function handleResetPrompt() {
+    eventSourceRef.current?.close()
+    eventSourceRef.current = undefined
+    setSearch('')
+    setAnswer(undefined)
+    setIsResponding(false)
+    setHasClippyError(false)
+  }
+
+  useEffect(() => {
+    if (search) {
+      handleConfirm(search)
+    }
+  }, [])
+
+  return (
+    <div onClick={(e) => e.stopPropagation()}>
+      <div className={cn('relative mb-[62px] py-4 max-h-[720px] overflow-auto')}>
+        {promptData.map((prompt, i) => {
+          if (!prompt.query) return <></>
+
+          return (
+            <>
+              {prompt.query && (
+                <div className="flex gap-6 mx-4 [overflow-anchor:none] mb-6">
+                  <div
+                    className="
+                      w-7 h-7 bg-scale-200 rounded-full border border-scale-400 flex items-center justify-center text-scale-1000 first-letter:
+                      ring-scale-200 ring-1 shadow-sm
+                  "
+                  >
+                    <IconUser strokeWidth={1.5} size={16} />
+                  </div>
+                  <div className="prose text-scale-1000">{prompt.query}</div>
+                </div>
+              )}
+
+              <div className="px-4 [overflow-anchor:none] mb-6">
+                {cantHelp ? (
+                  <p className="flex flex-col gap-4 items-center p-4">
+                    <div className="grid md:flex items-center gap-2 mt-4 text-center justify-items-center">
+                      <IconAlertCircle />
+                      <p>Sorry, I don&apos;t know how to help with that.</p>
+                    </div>
+                    <Button size="tiny" type="secondary" onClick={handleResetPrompt}>
+                      Try again?
+                    </Button>
+                  </p>
+                ) : (
+                  <div className="flex gap-6 [overflow-anchor:none] mb-6">
+                    <div>
+                      <AiIconChat />
+                    </div>
+                    <>
+                      {isLoading && promptIndex === i ? (
+                        <div className="bg-scale-700 h-[21px] w-[13px] mt-1 animate-pulse animate-bounce"></div>
+                      ) : (
+                        <div className="space-y-2 flex-grow">
+                          {/* @ts-expect-error */}
+                          <MarkdownHandler className="prose dark:prose-dark bg-scale-300 px-4 py-4 rounded-md w-full">
+                            {prompt.answer}
+                          </MarkdownHandler>
+                          {!isResponding && (
+                            <SQLOutputActions
+                              answer={prompt.answer}
+                              onSaveGeneratedSQL={onSaveGeneratedSQL}
+                            />
+                          )}
+                        </div>
+                      )}
+                    </>
+                  </div>
+                )}
+              </div>
+            </>
+          )
+        })}
+
+        {promptData.length === 0 && !hasClippyError && (
+          <CommandGroup heading="Examples" forceMount>
+            {SAMPLE_QUERIES.map((question) => {
+              const key = question.replace(/\s+/g, '_')
+              return (
+                <CommandItem
+                  type="command"
+                  onSelect={() => {
+                    if (!search) {
+                      handleConfirm(question)
+                    }
+                  }}
+                  forceMount
+                  key={key}
+                >
+                  <div>
+                    <AiIcon />
+                  </div>
+                  <p>{question}</p>
+                </CommandItem>
+              )
+            })}
+          </CommandGroup>
+        )}
+        {hasClippyError && (
+          <div className="p-6 flex flex-col items-center gap-6 mt-4">
+            <IconAlertTriangle className="text-amber-900" strokeWidth={1.5} size={21} />
+            <p className="text-lg text-scale-1200 text-center">
+              Sorry, looks like Clippy is having a hard time!
+            </p>
+            <p className="text-sm text-scale-900 text-center">Please try again in a bit.</p>
+            <Button size="tiny" type="secondary" onClick={handleResetPrompt}>
+              Try again?
+            </Button>
+          </div>
+        )}
+
+        <div className="[overflow-anchor:auto] h-px w-full"></div>
+      </div>
+      <div className="absolute bottom-0 w-full bg-scale-200 py-3">
+        <Input
+          className="bg-scale-100 rounded mx-3"
+          autoFocus
+          placeholder={
+            isLoading || isResponding
+              ? 'Waiting on an answer...'
+              : 'Describe what you need to Supabase AI, and it will try to generate the relevant SQL statements...'
+          }
+          value={search}
+          actions={
+            <>
+              {!isLoading && !isResponding ? (
+                <div
+                  className={`flex items-center gap-3 mr-3 transition-opacity duration-700 ${
+                    search ? 'opacity-100' : 'opacity-0'
+                  }`}
+                >
+                  <span className="text-scale-1100">Submit message</span>
+                  <div className="hidden text-scale-1100 md:flex items-center justify-center h-6 w-6 rounded bg-scale-500">
+                    <IconCornerDownLeft size={12} strokeWidth={1.5} />
+                  </div>
+                </div>
+              ) : null}
+            </>
+          }
+          onChange={(e) => {
+            if (!isLoading || !isResponding) {
+              setSearch(e.target.value)
+            }
+          }}
+          onKeyDown={(e) => {
+            switch (e.key) {
+              case 'Enter':
+                if (!search) {
+                  return
+                }
+                if (isLoading || isResponding) {
+                  return
+                }
+                handleConfirm(search)
+                return
+              default:
+                return
+            }
+          }}
+        />
+      </div>
+    </div>
+  )
+}
+
+export default GenerateSQL

studio/components/interfaces/App/CommandMenuWrapper.tsx

@@ -0,0 +1,55 @@
+import remarkGfm from 'remark-gfm'
+import ReactMarkdown from 'react-markdown'
+import { PropsWithChildren } from 'react'
+import { CommandMenuProvider } from 'ui'
+import { checkPermissions, useOptimisticSqlSnippetCreate, useParams, useStore } from 'hooks'
+import { PermissionAction } from '@supabase/shared-types/out/constants'
+import { useProfileQuery } from 'data/profile/profile-query'
+import { uuidv4 } from 'lib/helpers'
+import { createSqlSnippetSkeleton } from 'components/to-be-cleaned/SqlEditor/SqlEditor.utils'
+
+const CommandMenuWrapper = ({ children }: PropsWithChildren<{}>) => {
+  const { ref } = useParams()
+  const { content } = useStore()
+
+  const { data: profile } = useProfileQuery()
+  // const canCreateSQLSnippet = checkPermissions(PermissionAction.CREATE, 'user_content', {
+  //   resource: { type: 'sql', owner_id: profile?.id },
+  //   subject: { id: profile?.id },
+  // })
+
+  const onSaveGeneratedSQL = async (answer: string, resolve: any) => {
+    const formattedOutput = `
+-- Note: This query was generated via Supabase AI, please do verify the correctness of the
+-- SQL snippet before running it against your database as we are not able to guarantee the
+-- correctness of the snippet that was generated.
+
+${answer}
+`.trim()
+
+    console.log('onSaveGeneratedSQL', { formattedOutput })
+    const snippet = createSqlSnippetSkeleton({
+      owner_id: profile?.id,
+      name: 'Generated query',
+      sql: formattedOutput,
+    })
+    const payload = { id: uuidv4(), ...snippet }
+    await content.create(payload)
+    resolve()
+  }
+
+  return (
+    <CommandMenuProvider
+      site="studio"
+      projectRef={ref}
+      MarkdownHandler={({ ...props }) => (
+        <ReactMarkdown remarkPlugins={[remarkGfm]} components={null} {...props} />
+      )}
+      onSaveGeneratedSQL={onSaveGeneratedSQL}
+    >
+      {children}
+    </CommandMenuProvider>
+  )
+}
+
+export default CommandMenuWrapper

studio/components/interfaces/App/index.ts

@@ -1,5 +1,6 @@
 import PortalToast from './PortalToast'
 import RouteValidationWrapper from './RouteValidationWrapper'
 import AppBannerWrapper from './AppBannerWrapper'
+import CommandMenuWrapper from './CommandMenuWrapper'
 
-export { PortalToast, RouteValidationWrapper, AppBannerWrapper }
+export { PortalToast, RouteValidationWrapper, AppBannerWrapper, CommandMenuWrapper }

studio/pages/_app.tsx

@@ -31,25 +31,25 @@ import { Hydrate, QueryClientProvider } from '@tanstack/react-query'
 import { ReactQueryDevtools } from '@tanstack/react-query-devtools'
 import { RootStore } from 'stores'
 import HCaptchaLoadedStore from 'stores/hcaptcha-loaded-store'
-import { StoreProvider, useParams } from 'hooks'
+import { StoreProvider } from 'hooks'
 import { AuthProvider } from 'lib/auth'
 import { dart } from 'lib/constants/prism'
 import { useRootQueryClient } from 'data/query-client'
 
-import { PortalToast, RouteValidationWrapper, AppBannerWrapper } from 'components/interfaces/App'
+import {
+  PortalToast,
+  RouteValidationWrapper,
+  AppBannerWrapper,
+  CommandMenuWrapper,
+} from 'components/interfaces/App'
 import PageTelemetry from 'components/ui/PageTelemetry'
 import FlagProvider from 'components/ui/Flag/FlagProvider'
 import useAutoAuthRedirect from 'hooks/misc/useAutoAuthRedirect'
 
 import { TooltipProvider } from '@radix-ui/react-tooltip'
 import { IS_PLATFORM } from 'lib/constants'
-import { createBrowserSupabaseClient } from '@supabase/auth-helpers-nextjs'
 import { SessionContextProvider } from '@supabase/auth-helpers-react'
 import { createClient } from '@supabase/supabase-js'
-import { CommandMenuProvider } from 'ui'
-
-import remarkGfm from 'remark-gfm'
-import ReactMarkdown from 'react-markdown'
 
 dayjs.extend(customParseFormat)
 dayjs.extend(utc)
@@ -59,7 +59,6 @@ dayjs.extend(relativeTime)
 dart(Prism)
 
 function CustomApp({ Component, pageProps }: AppPropsWithLayout) {
-  const { ref } = useParams()
   const queryClient = useRootQueryClient()
   const [rootStore] = useState(() => new RootStore())
 
@@ -146,15 +145,9 @@ function CustomApp({ Component, pageProps }: AppPropsWithLayout) {
               <PageTelemetry>
                 <TooltipProvider>
                   <RouteValidationWrapper>
-                    <CommandMenuProvider
-                      site="studio"
-                      projectRef={ref}
-                      MarkdownHandler={({ ...props }) => (
-                        <ReactMarkdown remarkPlugins={[remarkGfm]} components={null} {...props} />
-                      )}
-                    >
+                    <CommandMenuWrapper>
                       <AppBannerWrapper>{getLayout(<Component {...pageProps} />)}</AppBannerWrapper>
-                    </CommandMenuProvider>
+                    </CommandMenuWrapper>
                   </RouteValidationWrapper>
                 </TooltipProvider>
               </PageTelemetry>