Add inline completions to SQL Editor (#30706)

* monaco complete

* progress

* just add panel to sqleditor

* update panel and prompt

* adjust prompt

* clean up

* clean up and improve prompt

* update prompt

* refine prompt

* remove package

* use constants file

* revert v3 prompt

* combine diff action bars

* use widget

* refactor

* Some updates

* fix: useCompletion base path

* fix: typo on monaco action label

* feat: use entire selected line(s) for inline completions

* fix widget consistency

* widget styling

* Fix TS

* feat: lock llm model version

* fix enter key issue

* move complete endpoint

* update complete prompt

* remove file

* placeholder text

---------

Co-authored-by: Joshen Lim <joshenlimek@gmail.com>
Co-authored-by: Greg Richardson <greg.nmr@gmail.com>

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

@@ -0,0 +1,79 @@
+import { detectOS } from 'lib/helpers'
+import { ArrowDownLeft, Loader2, Wand } from 'lucide-react'
+import { useCallback } from 'react'
+import { Input, Button } from 'ui'
+
+interface AskAIWidgetProps {
+  value: string
+  onChange: (value: string) => void
+  onSubmit: (prompt: string) => void
+  onAccept?: () => void
+  onReject?: () => void
+  isDiffVisible: boolean
+  isLoading?: boolean
+}
+
+export const AskAIWidget = ({
+  value,
+  onChange,
+  onSubmit,
+  onAccept,
+  onReject,
+  isDiffVisible,
+  isLoading = false,
+}: AskAIWidgetProps) => {
+  const os = detectOS()
+
+  // Auto focus input
+  const inputRef = useCallback((input: HTMLInputElement | null) => {
+    setTimeout(() => {
+      input?.focus()
+    }, 0)
+  }, [])
+
+  const handleSubmit = () => {
+    if (value.trim() && !isLoading) {
+      onSubmit(value)
+    }
+  }
+
+  return (
+    <div className="overflow-hidden rounded-md p-0 bg-popover border border-foreground/20 focus-within:border-foreground/30 shadow-xl text-sm max-w-xl">
+      <Input
+        inputRef={inputRef}
+        size="xlarge"
+        inputClassName="bg-transparent border-none shadow-none gap-4 text-xs focus-visible:outline-none focus-visible:ring-0 py-2 pl-3"
+        placeholder={isDiffVisible ? 'Make an edit...' : 'Edit SQL via the Assistant...'}
+        autoFocus
+        value={value}
+        onChange={(e) => onChange(e.target.value)}
+        onKeyDown={(e) => {
+          if (e.key === 'Enter' && !e.metaKey && !e.ctrlKey) {
+            handleSubmit()
+          }
+        }}
+        disabled={isLoading}
+      />
+      {isDiffVisible && (
+        <div className="flex justify-start p-0 border-t">
+          <Button
+            type="text"
+            onClick={onAccept}
+            className="text-xs h-auto py-1 rounded-none px-3 border-r-border"
+            disabled={isLoading}
+          >
+            Accept <span className="text-xs text-foreground-light">⌘ + Enter</span>
+          </Button>
+          <Button
+            onClick={onReject}
+            type="text"
+            className="text-xs h-auto py-1 rounded-none px-3 border-r-border"
+            disabled={isLoading}
+          >
+            Reject <span className="text-xs text-foreground-light">Esc</span>
+          </Button>
+        </div>
+      )}
+    </div>
+  )
+}

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

@@ -0,0 +1,108 @@
+import { editor } from 'monaco-editor'
+import { PropsWithChildren, useEffect, useMemo } from 'react'
+import { createPortal } from 'react-dom'
+
+export interface InlineWidgetProps {
+  /**
+   * The Monaco editor instance.
+   */
+  editor: editor.IStandaloneCodeEditor | editor.IStandaloneDiffEditor
+
+  /**
+   * ID used by Monaco to reference this widget
+   */
+  id: string
+
+  /**
+   * The line number after or before which this zone should appear.
+   * Use 0 to place a view zone before the first line number.
+   */
+  afterLineNumber: number
+  beforeLineNumber?: number
+
+  /**
+   * The height in lines of the view zone.
+   * @default 1
+   */
+  heightInLines?: number
+}
+
+/**
+ * Adds an inline widget to a Monaco editor instance. Acts as a
+ * container for custom widgets that are rendered inline.
+ *
+ * Implemented using the same techniques VS Code uses for their inline
+ * widgets (such as "Go to References"), but built for React.
+ * Uses a combination of a view zone and overlay widget.
+ */
+const InlineWidget = ({
+  children,
+  editor,
+  id,
+  beforeLineNumber,
+  afterLineNumber,
+  heightInLines = 1,
+}: PropsWithChildren<InlineWidgetProps>) => {
+  const lineNumber = beforeLineNumber ?? afterLineNumber
+  const key = `${id}-${lineNumber.toString()}`
+
+  const containerElement = useMemo(() => document.createElement('div'), [])
+
+  // Get the appropriate editor instance for diff editor
+  const targetEditor = 'getModifiedEditor' in editor ? editor.getModifiedEditor() : editor
+
+  useEffect(() => {
+    let zoneId: string
+    let viewZoneTop = 0
+    let viewZoneHeight = 0
+
+    const overlayWidget: editor.IOverlayWidget = {
+      getId: () => id,
+      getDomNode: () => containerElement,
+      getPosition: () => null,
+    }
+
+    const recalculateLayout = () => {
+      const layoutInfo = targetEditor.getLayoutInfo()
+
+      if (!layoutInfo) {
+        return
+      }
+
+      containerElement.style.left = `${layoutInfo.contentLeft}px`
+      containerElement.style.top = `${viewZoneTop}px`
+      containerElement.style.width = `${layoutInfo.width - layoutInfo.contentLeft}px`
+      containerElement.style.height = `${viewZoneHeight}px`
+    }
+
+    targetEditor.changeViewZones((accessor) => {
+      zoneId = accessor.addZone({
+        afterLineNumber: beforeLineNumber ?? afterLineNumber,
+        heightInLines,
+        domNode: document.createElement('div'),
+        onDomNodeTop: (top) => {
+          viewZoneTop = top
+          recalculateLayout()
+        },
+        onComputedHeight: (height) => {
+          viewZoneHeight = height
+          recalculateLayout()
+        },
+      })
+
+      targetEditor.addOverlayWidget(overlayWidget)
+    })
+
+    // Remove the view zone & overlay widget on unmount
+    return () => {
+      targetEditor.changeViewZones((accessor) => {
+        accessor.removeZone(zoneId)
+        targetEditor.removeOverlayWidget(overlayWidget)
+      })
+    }
+  }, [targetEditor, id, beforeLineNumber, afterLineNumber, heightInLines, containerElement])
+
+  return createPortal(children, containerElement, key)
+}
+
+export default InlineWidget

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

@@ -24,6 +24,13 @@ export type MonacoEditorProps = {
   autoFocus?: boolean
   executeQuery: () => void
   onHasSelection: (value: boolean) => void
+  onPrompt?: (value: {
+    selection: string
+    beforeSelection: string
+    afterSelection: string
+    startLineNumber: number
+    endLineNumber: number
+  }) => void
 }
 
 const MonacoEditor = ({
@@ -34,6 +41,7 @@ const MonacoEditor = ({
   className,
   executeQuery,
   onHasSelection,
+  onPrompt,
 }: MonacoEditorProps) => {
   const router = useRouter()
   const { profile } = useProfile()
@@ -92,6 +100,36 @@ const MonacoEditor = ({
       },
     })
 
+    if (onPrompt) {
+      editor.addAction({
+        id: 'generate-sql',
+        label: 'Generate SQL',
+        keybindings: [monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyK],
+        run: () => {
+          const selection = editor.getSelection()
+          const model = editor.getModel()
+          if (!model || !selection) return
+
+          const allLines = model.getLinesContent()
+
+          const startLineIndex = selection.startLineNumber - 1
+          const endLineIndex = selection.endLineNumber
+
+          const beforeSelection = allLines.slice(0, startLineIndex).join('\n') + '\n'
+          const selectedText = allLines.slice(startLineIndex, endLineIndex).join('\n')
+          const afterSelection = '\n' + allLines.slice(endLineIndex).join('\n')
+
+          onPrompt({
+            selection: selectedText,
+            beforeSelection,
+            afterSelection,
+            startLineNumber: selection?.startLineNumber ?? 0,
+            endLineNumber: selection?.endLineNumber ?? 0,
+          })
+        },
+      })
+    }
+
     editor.onDidChangeCursorSelection(({ selection }) => {
       const noSelection =
         selection.startLineNumber === selection.endLineNumber &&
@@ -99,25 +137,13 @@ const MonacoEditor = ({
       onHasSelection(!noSelection)
     })
 
-    // add margin above first line
-    editorRef.current.changeViewZones((accessor) => {
-      accessor.addZone({
-        afterLineNumber: 0,
-        heightInPx: 4,
-        domNode: document.createElement('div'),
-      })
-    })
-
     if (autoFocus) {
       if (editor.getValue().length === 1) editor.setPosition({ lineNumber: 1, column: 2 })
       editor.focus()
     }
   }
 
-  // [Joshen] Also needs updating here
-  const debouncedSetSql = debounce((id, value) => {
-    snapV2.setSql(id, value)
-  }, 1000)
+  const debouncedSetSql = debounce((id, value) => snapV2.setSql(id, value), 1000)
 
   function handleEditorChange(value: string | undefined) {
     const snippetCheck = snapV2.snippets[id]
@@ -145,9 +171,7 @@ const MonacoEditor = ({
   // if an SQL query is passed by the content parameter, set the editor value to its content. This
   // is usually used for sending the user to SQL editor from other pages with SQL.
   useEffect(() => {
-    if (content && content.length > 0) {
-      handleEditorChange(content)
-    }
+    if (content && content.length > 0) handleEditorChange(content)
   }, [])
 
   return (
@@ -171,9 +195,11 @@ const MonacoEditor = ({
         options={{
           tabSize: 2,
           fontSize: 13,
+          lineDecorationsWidth: 0,
           readOnly: disableEdit,
           minimap: { enabled: false },
           wordWrap: 'on',
+          padding: { top: 4 },
           // [Joshen] Commenting the following out as it causes the autocomplete suggestion popover
           // to be positioned wrongly somehow. I'm not sure if this affects anything though, but leaving
           // comment just in case anyone might be wondering. Relevant issues:

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

@@ -0,0 +1,66 @@
+import { Separator } from 'ui'
+import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
+
+interface RunQueryWarningModalProps {
+  visible: boolean
+  hasDestructiveOperations: boolean
+  hasUpdateWithoutWhere: boolean
+  onCancel: () => void
+  onConfirm: () => void
+}
+
+export const RunQueryWarningModal = ({
+  visible,
+  hasDestructiveOperations,
+  hasUpdateWithoutWhere,
+  onCancel,
+  onConfirm,
+}: RunQueryWarningModalProps) => {
+  return (
+    <ConfirmationModal
+      visible={visible}
+      size="large"
+      title={`Potential issue${hasDestructiveOperations && hasUpdateWithoutWhere ? 's' : ''} detected with your query`}
+      confirmLabel="Run this query"
+      variant="warning"
+      alert={{
+        base: {
+          variant: 'warning',
+        },
+        title:
+          hasDestructiveOperations && hasUpdateWithoutWhere
+            ? 'The following potential issues have been detected:'
+            : 'The following potential issue has been detected:',
+        description: 'Ensure that these are intentional before executing this query',
+      }}
+      onCancel={onCancel}
+      onConfirm={onConfirm}
+    >
+      <div className="text-sm">
+        <ul className="border rounded-md grid bg-surface-200">
+          {hasDestructiveOperations && (
+            <li className="grid pt-3 pb-2 px-4">
+              <span className="font-bold">Query has destructive operation</span>
+              <span className="text-foreground-lighter">
+                Make sure you are not accidentally removing something important.
+              </span>
+            </li>
+          )}
+          {hasDestructiveOperations && hasUpdateWithoutWhere && <Separator />}
+          {hasUpdateWithoutWhere && (
+            <li className="grid pt-2 pb-3 px-4 gap-1">
+              <span className="font-bold">Query uses update without a where clause</span>
+              <span className="text-foreground-lighter">
+                Without a <code className="text-xs">where</code> clause, this could update all rows
+                in the table.
+              </span>
+            </li>
+          )}
+        </ul>
+      </div>
+      <p className="mt-4 text-sm text-foreground-light">
+        Please confirm that you would like to execute this query.
+      </p>
+    </ConfirmationModal>
+  )
+}

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

@@ -1,18 +1,19 @@
 import type { Monaco } from '@monaco-editor/react'
 import { useQueryClient } from '@tanstack/react-query'
-import { motion } from 'framer-motion'
-import { ChevronUp, Loader2 } from 'lucide-react'
+import { AnimatePresence, motion } from 'framer-motion'
+import { ChevronUp, Command, Loader2 } from 'lucide-react'
 import dynamic from 'next/dynamic'
 import { useRouter } from 'next/router'
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import { toast } from 'sonner'
 
-import { Separator } from '@ui/components/SidePanel/SidePanel'
+import { useCompletion } from 'ai/react'
 import { useParams } from 'common'
 import { GridFooter } from 'components/ui/GridFooter'
 import { useSqlDebugMutation } from 'data/ai/sql-debug-mutation'
 import { useSqlTitleGenerateMutation } from 'data/ai/sql-title-mutation'
 import { useEntityDefinitionsQuery } from 'data/database/entity-definitions-query'
+import { constructHeaders } from 'data/fetchers'
 import { lintKeys } from 'data/lint/keys'
 import { useReadReplicasQuery } from 'data/read-replicas/replicas-query'
 import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
@@ -20,16 +21,16 @@ import { useFormatQueryMutation } from 'data/sql/format-sql-query'
 import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query'
 import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
 import { isError } from 'data/utils/error-check'
-import { useLocalStorageQuery } from 'hooks/misc/useLocalStorage'
 import { useOrgOptedIntoAi } from 'hooks/misc/useOrgOptedIntoAi'
 import { useSchemasForAi } from 'hooks/misc/useSchemasForAi'
 import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization'
 import { useSelectedProject } from 'hooks/misc/useSelectedProject'
-import { IS_PLATFORM, LOCAL_STORAGE_KEYS } from 'lib/constants'
+import { BASE_PATH, IS_PLATFORM, LOCAL_STORAGE_KEYS } from 'lib/constants'
+import { detectOS, uuidv4 } from 'lib/helpers'
 import { TelemetryActions } from 'lib/constants/telemetry'
-import { uuidv4 } from 'lib/helpers'
 import { useProfile } from 'lib/profile'
 import { wrapWithRoleImpersonation } from 'lib/role-impersonation'
+import { format } from 'sql-formatter'
 import { useAppStateSnapshot } from 'state/app-state'
 import { useDatabaseSelectorStateSnapshot } from 'state/database-selector'
 import { isRoleImpersonationEnabled, useGetImpersonatedRole } from 'state/role-impersonation-state'
@@ -47,28 +48,21 @@ import {
   TooltipContent_Shadcn_,
   TooltipTrigger_Shadcn_,
   Tooltip_Shadcn_,
-  cn,
 } from 'ui'
-import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
 import { subscriptionHasHipaaAddon } from '../Billing/Subscription/Subscription.utils'
-import { DiffActionBar } from './DiffActionBar'
+import { AskAIWidget } from './AskAIWidget'
+import { useSqlEditorDiff, useSqlEditorPrompt } from './hooks'
+import InlineWidget from './InlineWidget'
+import { RunQueryWarningModal } from './RunQueryWarningModal'
 import {
   ROWS_PER_PAGE_OPTIONS,
   sqlAiDisclaimerComment,
   untitledSnippetTitle,
 } from './SQLEditor.constants'
-import {
-  ContentDiff,
-  DiffType,
-  IStandaloneCodeEditor,
-  IStandaloneDiffEditor,
-} from './SQLEditor.types'
+import { DiffType, IStandaloneCodeEditor, IStandaloneDiffEditor } from './SQLEditor.types'
 import {
   checkDestructiveQuery,
   checkIfAppendLimitRequired,
-  compareAsAddition,
-  compareAsModification,
-  compareAsNewSnippet,
   createSqlSnippetSkeletonV2,
   isUpdateWithoutWhere,
   suffixWithLimit,
@@ -83,57 +77,69 @@ const DiffEditor = dynamic(
 )
 
 export const SQLEditor = () => {
+  const os = detectOS()
   const router = useRouter()
   const { ref, id: urlId } = useParams()
-
-  // generate an id to be used for new snippets. The dependency on urlId is to avoid a bug which
-  // shows up when clicking on the SQL Editor while being in the SQL editor on a random snippet.
-  const generatedId = useMemo(() => uuidv4(), [urlId])
-  // the id is stable across renders - it depends either on the url or on the memoized generated id
-  const id = !urlId || urlId === 'new' ? generatedId : urlId
-
   const { profile } = useProfile()
+  const queryClient = useQueryClient()
   const project = useSelectedProject()
   const organization = useSelectedOrganization()
   const appSnap = useAppStateSnapshot()
   const snapV2 = useSqlEditorV2StateSnapshot()
   const getImpersonatedRole = useGetImpersonatedRole()
   const databaseSelectorState = useDatabaseSelectorStateSnapshot()
-  const queryClient = useQueryClient()
+  const isOptedInToAI = useOrgOptedIntoAi()
+  const [selectedSchemas] = useSchemasForAi(project?.ref!)
+  const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM
 
-  const { mutate: formatQuery } = useFormatQueryMutation()
-  const { mutateAsync: generateSqlTitle } = useSqlTitleGenerateMutation()
-  const { mutateAsync: debugSql, isLoading: isDebugSqlLoading } = useSqlDebugMutation()
-
-  const [sourceSqlDiff, setSourceSqlDiff] = useState<ContentDiff>()
-  const [pendingTitle, setPendingTitle] = useState<string>()
-  const [hasSelection, setHasSelection] = useState<boolean>(false)
+  const {
+    sourceSqlDiff,
+    setSourceSqlDiff,
+    selectedDiffType,
+    setSelectedDiffType,
+    pendingTitle,
+    setPendingTitle,
+    isAcceptDiffLoading,
+    setIsAcceptDiffLoading,
+    isDiffOpen,
+    defaultSqlDiff,
+    closeDiff,
+  } = useSqlEditorDiff()
+  const { promptState, setPromptState, promptInput, setPromptInput, resetPrompt } =
+    useSqlEditorPrompt()
 
   const editorRef = useRef<IStandaloneCodeEditor | null>(null)
   const monacoRef = useRef<Monaco | null>(null)
   const diffEditorRef = useRef<IStandaloneDiffEditor | null>(null)
-
-  const { data: subscription } = useOrgSubscriptionQuery({ orgSlug: organization?.slug })
-  const { data: databases, isSuccess: isSuccessReadReplicas } = useReadReplicasQuery({
-    projectRef: ref,
-  })
-
+  const [hasSelection, setHasSelection] = useState<boolean>(false)
+  const [lineHighlights, setLineHighlights] = useState<string[]>([])
+  const [isDiffEditorMounted, setIsDiffEditorMounted] = useState(false)
   const [showPotentialIssuesModal, setShowPotentialIssuesModal] = useState(false)
   const [queryHasDestructiveOperations, setQueryHasDestructiveOperations] = useState(false)
   const [queryHasUpdateWithoutWhere, setQueryHasUpdateWithoutWhere] = useState(false)
+  const [showWidget, setShowWidget] = useState(false)
 
-  const isOptedInToAI = useOrgOptedIntoAi()
-  const [selectedSchemas] = useSchemasForAi(project?.ref!)
-  // Customers on HIPAA plans should not have access to Supabase AI
+  // generate an id to be used for new snippets. The dependency on urlId is to avoid a bug which
+  // shows up when clicking on the SQL Editor while being in the SQL editor on a random snippet.
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  const generatedId = useMemo(() => uuidv4(), [urlId])
+  // the id is stable across renders - it depends either on the url or on the memoized generated id
+  const id = !urlId || urlId === 'new' ? generatedId : urlId
+
+  const limit = snapV2.limit
+  const results = snapV2.results[id]?.[0]
+  const snippetIsLoading = !(
+    id in snapV2.snippets && snapV2.snippets[id].snippet.content !== undefined
+  )
+  const isLoading = urlId === 'new' ? false : snippetIsLoading
+
+  /** React query data fetching  */
+  const { data: subscription } = useOrgSubscriptionQuery({ orgSlug: organization?.slug })
   const hasHipaaAddon = subscriptionHasHipaaAddon(subscription)
-  const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM
 
-  const [isAcceptDiffLoading, setIsAcceptDiffLoading] = useState(false)
-  const [, setAiQueryCount] = useLocalStorageQuery('supabase_sql-editor-ai-query-count', 0)
-
-  const [selectedDiffType, setSelectedDiffType] = useState<DiffType | undefined>(undefined)
-  const [isFirstRender, setIsFirstRender] = useState(true)
-  const [lineHighlights, setLineHighlights] = useState<string[]>([])
+  const { data: databases, isSuccess: isSuccessReadReplicas } = useReadReplicasQuery({
+    projectRef: ref,
+  })
 
   const { data, refetch: refetchEntityDefinitions } = useEntityDefinitionsQuery(
     {
@@ -143,19 +149,13 @@ export const SQLEditor = () => {
     },
     { enabled: includeSchemaMetadata }
   )
-
   const entityDefinitions = includeSchemaMetadata ? data?.map((def) => def.sql.trim()) : undefined
-  const isDiffOpen = !!sourceSqlDiff
-
-  const limit = snapV2.limit
-  const results = snapV2.results[id]?.[0]
-  const snippetIsLoading = !(
-    id in snapV2.snippets && snapV2.snippets[id].snippet.content !== undefined
-  )
-  const isLoading = urlId === 'new' ? false : snippetIsLoading
 
+  /* React query mutations */
+  const { mutate: formatQuery } = useFormatQueryMutation()
+  const { mutateAsync: generateSqlTitle } = useSqlTitleGenerateMutation()
+  const { mutateAsync: debugSql, isLoading: isDebugSqlLoading } = useSqlDebugMutation()
   const { mutate: sendEvent } = useSendEventMutation()
-
   const { mutate: execute, isLoading: isExecuting } = useExecuteSqlMutation({
     onSuccess(data, vars) {
       if (id) snapV2.addResult(id, data.result, vars.autoLimit)
@@ -327,6 +327,7 @@ export const SQLEditor = () => {
         })
       }
     },
+    // eslint-disable-next-line react-hooks/exhaustive-deps
     [
       isDiffOpen,
       id,
@@ -363,35 +364,10 @@ export const SQLEditor = () => {
         toast.error(`Failed to create new query: ${error.message}`)
       }
     },
+    // eslint-disable-next-line react-hooks/exhaustive-deps
     [profile?.id, project?.id, ref, router, snapV2]
   )
 
-  const updateEditorWithCheckForDiff = useCallback(
-    ({ diffType, sql }: { diffType: DiffType; sql: string }) => {
-      const editorModel = editorRef.current?.getModel()
-      if (!editorModel) return
-
-      setAiQueryCount((count) => count + 1)
-
-      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()
-        const diff = { original: currentSql || '', modified: sql }
-        setSourceSqlDiff(diff)
-        setSelectedDiffType(diffType)
-      }
-    },
-    [setAiQueryCount]
-  )
-
   const onDebug = useCallback(async () => {
     try {
       const snippet = snapV2.snippets[id]
@@ -411,27 +387,20 @@ export const SQLEditor = () => {
         )
       }
     }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [debugSql, entityDefinitions, id, snapV2.results, snapV2.snippets])
 
   const acceptAiHandler = useCallback(async () => {
     try {
       setIsAcceptDiffLoading(true)
 
-      if (!sourceSqlDiff) {
-        return
-      }
-
       // TODO: show error if undefined
-      if (!editorRef.current || !diffEditorRef.current) {
-        return
-      }
+      if (!sourceSqlDiff || !editorRef.current || !diffEditorRef.current) return
 
       const editorModel = editorRef.current.getModel()
       const diffModel = diffEditorRef.current.getModel()
 
-      if (!editorModel || !diffModel) {
-        return
-      }
+      if (!editorModel || !diffModel) return
 
       const sql = diffModel.modified.getValue()
 
@@ -454,11 +423,12 @@ export const SQLEditor = () => {
       sendEvent({ action: TelemetryActions.ASSISTANT_SUGGESTION_ACCEPTED })
 
       setSelectedDiffType(DiffType.Modification)
-      setSourceSqlDiff(undefined)
-      setPendingTitle(undefined)
+      resetPrompt()
+      closeDiff()
     } finally {
       setIsAcceptDiffLoading(false)
     }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [
     sourceSqlDiff,
     selectedDiffType,
@@ -472,129 +442,201 @@ export const SQLEditor = () => {
 
   const discardAiHandler = useCallback(() => {
     sendEvent({ action: TelemetryActions.ASSISTANT_SUGGESTION_REJECTED })
-    setSourceSqlDiff(undefined)
-    setPendingTitle(undefined)
-  }, [router])
+    resetPrompt()
+    closeDiff()
+  }, [closeDiff, resetPrompt, sendEvent])
+
+  const {
+    complete,
+    completion,
+    isLoading: isCompletionLoading,
+  } = useCompletion({
+    api: `${BASE_PATH}/api/ai/sql/complete`,
+    body: {
+      projectRef: project?.ref,
+      connectionString: project?.connectionString,
+      includeSchemaMetadata,
+    },
+    onResponse: (response) => {
+      if (!response.ok) throw new Error('Failed to generate completion')
+    },
+    onError: (error) => {
+      toast.error(`Failed to generate SQL: ${error.message}`)
+    },
+  })
+
+  const handlePrompt = async (
+    prompt: string,
+    context: {
+      beforeSelection: string
+      selection: string
+      afterSelection: string
+    }
+  ) => {
+    try {
+      setPromptState((prev) => ({
+        ...prev,
+        selection: context.selection,
+        beforeSelection: context.beforeSelection,
+        afterSelection: context.afterSelection,
+      }))
+      const headerData = await constructHeaders()
+
+      await complete(prompt, {
+        headers: { Authorization: headerData.get('Authorization') ?? '' },
+        body: {
+          completionMetadata: {
+            textBeforeCursor: context.beforeSelection,
+            textAfterCursor: context.afterSelection,
+            language: 'pgsql',
+            prompt,
+            selection: context.selection,
+          },
+        },
+      })
+    } catch (error) {
+      setPromptState((prev) => ({ ...prev, isLoading: false }))
+    }
+  }
+
+  /** All useEffects are at the bottom before returning the TSX */
+
+  useEffect(() => {
+    if (id) {
+      closeDiff()
+      setPromptState((prev) => ({ ...prev, isOpen: false }))
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [closeDiff, id])
 
   useEffect(() => {
     const handler = (e: KeyboardEvent) => {
-      if (!isDiffOpen) {
-        return
-      }
+      if (!isDiffOpen && !promptState.isOpen) return
 
       switch (e.key) {
         case 'Enter':
-          acceptAiHandler()
+          if ((os === 'macos' ? e.metaKey : e.ctrlKey) && isDiffOpen) {
+            acceptAiHandler()
+            resetPrompt()
+          }
           return
         case 'Escape':
-          discardAiHandler()
+          if (isDiffOpen) discardAiHandler()
+          resetPrompt()
+          editorRef.current?.focus()
           return
       }
     }
-
     window.addEventListener('keydown', handler)
-
     return () => window.removeEventListener('keydown', handler)
-  }, [isDiffOpen, acceptAiHandler, discardAiHandler])
+  }, [os, isDiffOpen, promptState.isOpen, acceptAiHandler, discardAiHandler, resetPrompt])
 
   useEffect(() => {
-    const applyDiff = ({ original, modified }: { original: string; modified: string }) => {
-      const model = diffEditorRef.current?.getModel()
+    if (isDiffOpen) {
+      const diffEditor = diffEditorRef.current
+      const model = diffEditor?.getModel()
       if (model && model.original && model.modified) {
-        model.original.setValue(original)
-        model.modified.setValue(modified)
+        model.original.setValue(defaultSqlDiff.original)
+        model.modified.setValue(defaultSqlDiff.modified)
+        // scroll to the start line of the modification
+        const modifiedEditor = diffEditor!.getModifiedEditor()
+        const startLine = promptState.startLineNumber
+        modifiedEditor.revealLineInCenter(startLine)
       }
     }
-
-    const model = diffEditorRef.current?.getModel()
-    try {
-      if (model?.original && model.modified && sourceSqlDiff) {
-        switch (selectedDiffType) {
-          case DiffType.Modification: {
-            const transformedDiff = compareAsModification(sourceSqlDiff)
-            applyDiff(transformedDiff)
-            return
-          }
-
-          case DiffType.Addition: {
-            const transformedDiff = compareAsAddition(sourceSqlDiff)
-            applyDiff(transformedDiff)
-            return
-          }
-
-          case DiffType.NewSnippet: {
-            const transformedDiff = compareAsNewSnippet(sourceSqlDiff)
-            applyDiff(transformedDiff)
-            return
-          }
-
-          default:
-            throw new Error(`Unknown diff type '${selectedDiffType}'`)
-        }
-      }
-    } catch (e) {
-      console.log(e)
-    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [selectedDiffType, sourceSqlDiff])
 
-  // Used for cleaner framer motion transitions
-  useEffect(() => {
-    setIsFirstRender(false)
-  }, [])
-
   useEffect(() => {
     if (isSuccessReadReplicas) {
       const primaryDatabase = databases.find((db) => db.identifier === ref)
       databaseSelectorState.setSelectedDatabaseId(primaryDatabase?.identifier)
     }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [isSuccessReadReplicas, databases, ref])
 
   useEffect(() => {
     if (snapV2.diffContent !== undefined) {
-      updateEditorWithCheckForDiff(snapV2.diffContent)
+      const { diffType, sql }: { diffType: DiffType; sql: string } = snapV2.diffContent
+      const editorModel = editorRef.current?.getModel()
+      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: `${sql}`,
+            range: editorModel.getFullModelRange(),
+          },
+        ])
+      } else {
+        const currentSql = editorRef.current?.getValue()
+        const diff = { original: currentSql || '', modified: sql }
+        setSourceSqlDiff(diff)
+        setSelectedDiffType(diffType)
+      }
     }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [snapV2.diffContent])
 
-  const defaultSqlDiff = useMemo(() => {
-    if (!sourceSqlDiff) {
-      return { original: '', modified: '' }
+  useEffect(() => {
+    if (!completion) {
+      return
     }
-    switch (selectedDiffType) {
-      case DiffType.Modification: {
-        return compareAsModification(sourceSqlDiff)
-      }
 
-      case DiffType.Addition: {
-        return compareAsAddition(sourceSqlDiff)
-      }
+    const original =
+      promptState.beforeSelection + promptState.selection + promptState.afterSelection
+    const modified = promptState.beforeSelection + completion + promptState.afterSelection
 
-      case DiffType.NewSnippet: {
-        return compareAsNewSnippet(sourceSqlDiff)
-      }
+    if (isCompletionLoading) {
+      let formattedModified = modified
 
-      default:
-        return { original: '', modified: '' }
+      // Attempt to format the modified SQL in case the LLM left out indentation, etc
+      try {
+        formattedModified = format(
+          promptState.beforeSelection + completion + promptState.afterSelection,
+          {
+            language: 'postgresql',
+            keywordCase: 'lower',
+          }
+        )
+      } catch (error) {}
+
+      setSourceSqlDiff({
+        original,
+        modified: formattedModified,
+      })
+      setSelectedDiffType(DiffType.Modification)
+      setPromptState((prev) => ({ ...prev, isLoading: false }))
     }
-  }, [selectedDiffType, sourceSqlDiff])
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [
+    completion,
+    promptState.beforeSelection,
+    promptState.selection,
+    promptState.afterSelection,
+    isCompletionLoading,
+  ])
+
+  // We want to check if the diff editor is mounted and if it is, we want to show the widget
+  // We also want to cleanup the widget when the diff editor is closed
+  useEffect(() => {
+    if (!isDiffOpen) {
+      setIsDiffEditorMounted(false)
+      setShowWidget(false)
+    } else if (diffEditorRef.current && isDiffEditorMounted) {
+      setShowWidget(true)
+      return () => setShowWidget(false)
+    }
+  }, [isDiffOpen, isDiffEditorMounted])
 
   return (
     <>
-      <ConfirmationModal
+      <RunQueryWarningModal
         visible={showPotentialIssuesModal}
-        size="large"
-        title={`Potential issue${queryHasDestructiveOperations && queryHasUpdateWithoutWhere ? 's' : ''} detected with your query`}
-        confirmLabel="Run this query"
-        variant="warning"
-        alert={{
-          base: {
-            variant: 'warning',
-          },
-          title:
-            queryHasDestructiveOperations && queryHasUpdateWithoutWhere
-              ? 'The following potential issues have been detected:'
-              : 'The following potential issue has been detected:',
-          description: 'Ensure that these are intentional before executing this query',
-        }}
+        hasDestructiveOperations={queryHasDestructiveOperations}
+        hasUpdateWithoutWhere={queryHasUpdateWithoutWhere}
         onCancel={() => {
           setShowPotentialIssuesModal(false)
           setQueryHasDestructiveOperations(false)
@@ -605,33 +647,7 @@ export const SQLEditor = () => {
           setShowPotentialIssuesModal(false)
           executeQuery(true)
         }}
-      >
-        <div className="text-sm">
-          <ul className="border rounded-md grid bg-surface-200">
-            {queryHasDestructiveOperations && (
-              <li className="grid pt-3 pb-2 px-4">
-                <span className="font-bold">Query has destructive operation</span>
-                <span className="text-foreground-lighter">
-                  Make sure you are not accidentally removing something important.
-                </span>
-              </li>
-            )}
-            {queryHasDestructiveOperations && queryHasUpdateWithoutWhere && <Separator />}
-            {queryHasUpdateWithoutWhere && (
-              <li className="grid pt-2 pb-3 px-4 gap-1">
-                <span className="font-bold">Query uses update without a where clause</span>
-                <span className="text-foreground-lighter">
-                  Without a <code className="text-xs">where</code> clause, this could update all
-                  rows in the table.
-                </span>
-              </li>
-            )}
-          </ul>
-        </div>
-        <p className="mt-4 text-sm text-foreground-light">
-          Please confirm that you would like to execute this query.
-        </p>
-      </ConfirmationModal>
+      />
 
       <ResizablePanelGroup
         className="flex h-full"
@@ -644,27 +660,6 @@ export const SQLEditor = () => {
             direction="vertical"
             autoSaveId={LOCAL_STORAGE_KEYS.SQL_EDITOR_SPLIT_SIZE}
           >
-            {!hasHipaaAddon && isDiffOpen && (
-              <motion.div
-                key="ask-ai-input-container"
-                layoutId="ask-ai-input-container"
-                variants={{ visible: { borderRadius: 0, x: 0 }, hidden: { x: 100 } }}
-                initial={isFirstRender ? 'visible' : 'hidden'}
-                animate="visible"
-                className={cn(
-                  'flex flex-row items-center gap-3 justify-end px-2 py-2 w-full z-10',
-                  'bg-brand-200 border-b border-brand-400  !shadow-none'
-                )}
-              >
-                <DiffActionBar
-                  loading={isAcceptDiffLoading}
-                  selectedDiffType={selectedDiffType || DiffType.Modification}
-                  onChangeDiffType={(diffType) => setSelectedDiffType(diffType)}
-                  onAccept={acceptAiHandler}
-                  onCancel={discardAiHandler}
-                />
-              </motion.div>
-            )}
             <ResizablePanel maxSize={70}>
               <div className="flex-grow overflow-y-auto border-b h-full">
                 {isLoading ? (
@@ -674,15 +669,7 @@ export const SQLEditor = () => {
                 ) : (
                   <>
                     {isDiffOpen && (
-                      <motion.div
-                        className="w-full h-full"
-                        variants={{
-                          visible: { opacity: 1, filter: 'blur(0px)' },
-                          hidden: { opacity: 0, filter: 'blur(10px)' },
-                        }}
-                        initial="hidden"
-                        animate="visible"
-                      >
+                      <div className="w-full h-full">
                         <DiffEditor
                           theme="supabase"
                           language="pgsql"
@@ -690,21 +677,47 @@ export const SQLEditor = () => {
                           modified={defaultSqlDiff.modified}
                           onMount={(editor) => {
                             diffEditorRef.current = editor
+                            setIsDiffEditorMounted(true)
+                          }}
+                          options={{
+                            fontSize: 13,
+                            renderSideBySide: false,
+                            minimap: { enabled: false },
+                            wordWrap: 'on',
+                            lineNumbers: 'on',
+                            folding: false,
+                            padding: { top: 4 },
+                            lineNumbersMinChars: 3,
                           }}
-                          options={{ fontSize: 13 }}
                         />
-                      </motion.div>
+                        {showWidget && (
+                          <InlineWidget
+                            editor={diffEditorRef.current!}
+                            id="ask-ai-diff"
+                            heightInLines={3}
+                            afterLineNumber={0}
+                            beforeLineNumber={Math.max(0, promptState.startLineNumber - 1)}
+                          >
+                            <AskAIWidget
+                              onSubmit={(prompt: string) => {
+                                handlePrompt(prompt, {
+                                  beforeSelection: promptState.beforeSelection,
+                                  selection: promptState.selection || defaultSqlDiff.modified,
+                                  afterSelection: promptState.afterSelection,
+                                })
+                              }}
+                              value={promptInput}
+                              onChange={setPromptInput}
+                              onAccept={acceptAiHandler}
+                              onReject={discardAiHandler}
+                              isDiffVisible={true}
+                              isLoading={isCompletionLoading}
+                            />
+                          </InlineWidget>
+                        )}
+                      </div>
                     )}
-                    <motion.div
-                      key={id}
-                      variants={{
-                        visible: { opacity: 1, filter: 'blur(0px)' },
-                        hidden: { opacity: 0, filter: 'blur(10px)' },
-                      }}
-                      initial="hidden"
-                      animate={isDiffOpen ? 'hidden' : 'visible'}
-                      className="w-full h-full"
-                    >
+                    <div key={id} className="w-full h-full relative">
                       <MonacoEditor
                         autoFocus
                         id={id}
@@ -712,8 +725,61 @@ export const SQLEditor = () => {
                         monacoRef={monacoRef}
                         executeQuery={executeQuery}
                         onHasSelection={setHasSelection}
+                        onPrompt={({
+                          selection,
+                          beforeSelection,
+                          afterSelection,
+                          startLineNumber,
+                          endLineNumber,
+                        }) => {
+                          setPromptState((prev) => ({
+                            ...prev,
+                            isOpen: true,
+                            selection,
+                            beforeSelection,
+                            afterSelection,
+                            startLineNumber,
+                            endLineNumber,
+                          }))
+                        }}
                       />
-                    </motion.div>
+                      {editorRef.current && promptState.isOpen && !isDiffOpen && (
+                        <InlineWidget
+                          editor={editorRef.current}
+                          id="ask-ai"
+                          afterLineNumber={promptState.endLineNumber}
+                          beforeLineNumber={Math.max(0, promptState.startLineNumber - 1)}
+                          heightInLines={2}
+                        >
+                          <AskAIWidget
+                            value={promptInput}
+                            onChange={setPromptInput}
+                            onSubmit={(prompt: string) => {
+                              handlePrompt(prompt, {
+                                beforeSelection: promptState.beforeSelection,
+                                selection: promptState.selection,
+                                afterSelection: promptState.afterSelection,
+                              })
+                            }}
+                            isDiffVisible={false}
+                            isLoading={isCompletionLoading}
+                          />
+                        </InlineWidget>
+                      )}
+                      <AnimatePresence>
+                        {!promptState.isOpen && !editorRef.current?.getValue() && (
+                          <motion.p
+                            initial={{ y: 5, opacity: 0 }}
+                            animate={{ y: 0, opacity: 1 }}
+                            exit={{ y: 5, opacity: 0 }}
+                            className="text-foreground-lighter absolute bottom-4 left-4 z-10 font-mono text-xs flex items-center gap-1"
+                          >
+                            Hit {os === 'macos' ? <Command size={12} /> : `CTRL+`}K to edit with the
+                            Assistant
+                          </motion.p>
+                        )}
+                      </AnimatePresence>
+                    </div>
                   </>
                 )}
               </div>

apps/studio/components/interfaces/SQLEditor/SQLEditor.utils.ts

@@ -137,18 +137,14 @@ export const compareAsModification = (sqlDiff: ContentDiff) => {
 
   return {
     original: sqlDiff.original,
-    modified: `${sqlAiDisclaimerComment}\n\n${formattedModified}`,
+    modified: `${formattedModified}`,
   }
 }
 
 export const compareAsAddition = (sqlDiff: ContentDiff) => {
   const formattedOriginal = sqlDiff.original.replace(sqlAiDisclaimerComment, '').trim()
   const formattedModified = sqlDiff.modified.replace(sqlAiDisclaimerComment, '').trim()
-  const newModified =
-    sqlAiDisclaimerComment +
-    '\n\n' +
-    (formattedOriginal ? formattedOriginal + '\n\n' : '') +
-    formattedModified
+  const newModified = (formattedOriginal ? formattedOriginal + '\n\n' : '') + formattedModified
 
   return {
     original: sqlDiff.original,

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

@@ -113,7 +113,7 @@ const UtilityPanel = ({
 
   return (
     <Tabs_Shadcn_ defaultValue="results" className="w-full h-full flex flex-col">
-      <TabsList_Shadcn_ className="flex justify-between gap-2 px-2 overflow-x-auto min-h-[42px]">
+      <TabsList_Shadcn_ className="flex justify-between gap-2 px-4 overflow-x-auto min-h-[42px]">
         <div className="flex items-center gap-4">
           <TabsTrigger_Shadcn_ className="py-3 text-xs" value="results">
             <span className="translate-y-[1px]">Results</span>

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

@@ -156,7 +156,7 @@ const UtilityTabResults = ({
   } else if (!result) {
     return (
       <div className="bg-table-header-light [[data-theme*=dark]_&]:bg-table-header-dark">
-        <p className="m-0 border-0 px-6 py-4 text-sm text-foreground-light">
+        <p className="m-0 border-0 px-4 py-4 text-sm text-foreground-light">
           Click <code>Run</code> to execute your query.
         </p>
       </div>

apps/studio/components/interfaces/SQLEditor/hooks.ts

@@ -1,5 +1,6 @@
 import { PermissionAction } from '@supabase/shared-types/out/constants'
 import { useRouter } from 'next/router'
+import { useCallback, useEffect, useMemo, useState } from 'react'
 import { toast } from 'sonner'
 
 import { useParams } from 'common'
@@ -8,7 +9,13 @@ import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
 import { uuidv4 } from 'lib/helpers'
 import { useProfile } from 'lib/profile'
 import { useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2'
-import { createSqlSnippetSkeletonV2 } from './SQLEditor.utils'
+import { ContentDiff, DiffType } from './SQLEditor.types'
+import {
+  compareAsAddition,
+  compareAsModification,
+  compareAsNewSnippet,
+  createSqlSnippetSkeletonV2,
+} from './SQLEditor.utils'
 
 export const useNewQuery = () => {
   const router = useRouter()
@@ -50,4 +57,92 @@ export const useNewQuery = () => {
   return { newQuery }
 }
 
+export function useSqlEditorDiff() {
+  const [sourceSqlDiff, setSourceSqlDiff] = useState<ContentDiff>()
+  const [selectedDiffType, setSelectedDiffType] = useState<DiffType>()
+  const [pendingTitle, setPendingTitle] = useState<string>()
+  const [isAcceptDiffLoading, setIsAcceptDiffLoading] = useState(false)
+
+  const isDiffOpen = !!sourceSqlDiff
+
+  const defaultSqlDiff = useMemo(() => {
+    if (!sourceSqlDiff) {
+      return { original: '', modified: '' }
+    }
+
+    switch (selectedDiffType) {
+      case DiffType.Modification:
+        return compareAsModification(sourceSqlDiff)
+      case DiffType.Addition:
+        return compareAsAddition(sourceSqlDiff)
+      case DiffType.NewSnippet:
+        return compareAsNewSnippet(sourceSqlDiff)
+      default:
+        return { original: '', modified: '' }
+    }
+  }, [selectedDiffType, sourceSqlDiff])
+
+  const closeDiff = useCallback(() => {
+    setSourceSqlDiff(undefined)
+    setPendingTitle(undefined)
+    setSelectedDiffType(undefined)
+  }, [])
+
+  return {
+    sourceSqlDiff,
+    setSourceSqlDiff,
+    selectedDiffType,
+    setSelectedDiffType,
+    pendingTitle,
+    setPendingTitle,
+    isAcceptDiffLoading,
+    setIsAcceptDiffLoading,
+    isDiffOpen,
+    defaultSqlDiff,
+    closeDiff,
+  }
+}
+
+interface PromptState {
+  isOpen: boolean
+  selection: string
+  beforeSelection: string
+  afterSelection: string
+  startLineNumber: number
+  endLineNumber: number
+}
+
+const initialPromptState: PromptState = {
+  isOpen: false,
+  selection: '',
+  beforeSelection: '',
+  afterSelection: '',
+  startLineNumber: 0,
+  endLineNumber: 0,
+}
+
+export function useSqlEditorPrompt() {
+  const [promptState, setPromptState] = useState<PromptState>(initialPromptState)
+  const [promptInput, setPromptInput] = useState('')
+
+  useEffect(() => {
+    if (!promptState.isOpen) {
+      setPromptInput('')
+    }
+  }, [promptState.isOpen])
+
+  const resetPrompt = () => {
+    setPromptState(initialPromptState)
+    setPromptInput('')
+  }
+
+  return {
+    promptState,
+    setPromptState,
+    promptInput,
+    setPromptInput,
+    resetPrompt,
+  }
+}
+
 export default useNewQuery

apps/studio/middleware.ts

@@ -10,6 +10,7 @@ const HOSTED_SUPPORTED_API_URLS = [
   '/ai/sql/suggest',
   '/ai/sql/generate-v2',
   '/ai/sql/generate-v3',
+  '/ai/sql/complete',
   '/ai/sql/title',
   '/ai/sql/debug',
   '/ai/sql/cron',

apps/studio/package.json

@@ -30,6 +30,7 @@
     "@hookform/resolvers": "^3.1.1",
     "@monaco-editor/react": "^4.6.0",
     "@next/bundle-analyzer": "^14.2.3",
+    "@number-flow/react": "^0.3.2",
     "@radix-ui/react-popover": "^1.0.7",
     "@radix-ui/react-tooltip": "^1.0.7",
     "@sentry/nextjs": "^8.34.0",
@@ -77,7 +78,6 @@
     "mobx": "^6.10.2",
     "mobx-react-lite": "^4.0.5",
     "monaco-editor": "0.33.0",
-    "@number-flow/react": "^0.3.2",
     "next": "^14.2.10",
     "next-themes": "^0.3.0",
     "nuqs": "^1.17.8",

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

@@ -0,0 +1,148 @@
+import { streamText } from 'ai'
+import { openai } from '@ai-sdk/openai'
+import { NextApiRequest, NextApiResponse } from 'next'
+import { getTools } from '../sql/tools'
+import { executeSql } from 'data/sql/execute-sql-query'
+import pgMeta from '@supabase/pg-meta'
+
+export const maxDuration = 30
+const openAiKey = process.env.OPENAI_API_KEY
+const pgMetaSchemasList = pgMeta.schemas.list()
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse) {
+  if (!openAiKey) {
+    return new Response(
+      JSON.stringify({
+        error: 'No OPENAI_API_KEY set. Create this environment variable to use AI features.',
+      }),
+      {
+        status: 500,
+        headers: { 'Content-Type': 'application/json' },
+      }
+    )
+  }
+
+  if (req.method !== 'POST') {
+    return new Response(
+      JSON.stringify({ data: null, error: { message: `Method ${req.method} Not Allowed` } }),
+      {
+        status: 405,
+        headers: { 'Content-Type': 'application/json', Allow: 'POST' },
+      }
+    )
+  }
+
+  try {
+    const { completionMetadata, projectRef, connectionString, includeSchemaMetadata } = req.body
+    const { textBeforeCursor, textAfterCursor, language, prompt, selection } = completionMetadata
+
+    if (!projectRef) {
+      return res.status(400).json({
+        error: 'Missing project_ref in request body',
+      })
+    }
+
+    const authorization = req.headers.authorization
+
+    const { result: schemas } = includeSchemaMetadata
+      ? await executeSql(
+          {
+            projectRef,
+            connectionString,
+            sql: pgMetaSchemasList.sql,
+          },
+          undefined,
+          {
+            'Content-Type': 'application/json',
+            ...(authorization && { Authorization: authorization }),
+          }
+        )
+      : { result: [] }
+
+    const result = await streamText({
+      model: openai('gpt-4o-mini-2024-07-18'),
+      maxSteps: 5,
+      tools: getTools({ projectRef, connectionString, authorization, includeSchemaMetadata }),
+      system: `
+      You are a Supabase Postgres expert who can do the following things.
+
+      # You generate and debug SQL
+      The generated SQL (must be valid SQL), and must adhere to the following:
+      - Always use double apostrophe in SQL strings (eg. 'Night''s watch')
+      - Always use semicolons
+      - Use vector(384) data type for any embedding/vector related query
+      - When debugging, retrieve sql schema details to ensure sql is correct
+      - In Supabase, the auth schema already has a users table which is used to store users. It is common practice to create a profiles table in the public schema that links to auth.users to store user information instead. You don't need to create a new users table.
+      - Never suggest creating a view to retrieve information from the users table of the auth schema. This is against our best practices.
+
+      When generating tables, do the following:
+      - Ensure that all tables always have a primary key
+      - Ensure that all tables have RLS enabled. Inform the user that they will need to create RLS policies before being able to read or write to the table over Supabase APIs.
+      - For primary keys, always use "id bigint primary key generated always as identity" (not serial)
+      - Prefer creating foreign key references in the create statement
+      - Prefer 'text' over 'varchar'
+      - Prefer 'timestamp with time zone' over 'date'
+      - If a foreign key is created, include a statement to create a database index on the columns involved in the foreign key. This is a best practice as it is common to use foreign key columns in join conditions when querying the database, hence the index will improve the performance of those joins and reduce database resource consumption.
+      - If the table is a foreign table, create it in a new schema called private. If the private schema doesnt exist, create it. 
+        - Explain to the user that we create it outside of the public schema as it can present a security risk as foreign tables views cannot be configured to respect RLS policies. Provide this link for users to view more details about the security risk: https://supabase.com/docs/guides/database/database-advisors?queryGroups=lint&lint=0017_foreign_table_in_api
+
+      When generating views, do the following:
+      - All views should include 'with (security_invoker=on)' clause in the SQL statement for creating views.
+      - Place the 'with (security_invoker=on)' immediately after the CREATE VIEW statement, before AS
+      - If the view is a materialized view, create it in a new schema called private. If the private schema doesnt exist, create it. 
+        - Explain to the user that we create it outside of the public schema as it can present a security risk as materialized views cannot be configured to respect RLS policies of the underlying tables they are built upon, nor can they be secured with RLS directly. Provide this link for users to view more details about the security risk: https://supabase.com/docs/guides/database/database-advisors?queryGroups=lint&lint=0016_materialized_view_in_api
+
+      Feel free to suggest corrections for suspected typos.
+
+      # You write row level security policies.
+
+      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.
+      - When writing policies that use a function from the auth schema, ensure that the calls are wrapped with parentheses e.g select auth.uid() should be written as (select auth.uid()) instead
+
+      # 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
+      if the function returns a type of trigger. When generating functions, do the following:
+      - If the function returns a trigger type, ensure that it uses security definer, otherwise default to security invoker. Include this in the create functions SQL statement.
+      - Ensure to set the search_path configuration parameter as '', include this in the create functions SQL statement.
+      - Default to create or replace whenever possible for updating an existing function, otherwise use the alter function statement
+      Please make sure that all queries are valid Postgres SQL queries
+
+      # For all your abilities, 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.
+      - If the question is about users or involves creating a users table, also retrieve the auth schema.  
+
+      Here are the existing database schema names you can retrieve: ${schemas}
+      `,
+      messages: [
+        {
+          role: 'user',
+          content: `You are helping me edit some pgsql code.
+            Here is the context:
+            ${textBeforeCursor}<selection>${selection}</selection>${textAfterCursor}
+            
+            Instructions:
+            1. Only modify the selected text based on this prompt: ${prompt}
+            2. Your response should be ONLY the modified selection text, nothing else. Remove selected text if needed.
+            3. Do not wrap in code blocks or markdown
+            4. You can respond with one word or multiple words
+            5. Ensure the modified text flows naturally within the current line
+            6. Avoid duplicating SQL keywords (SELECT, FROM, WHERE, etc) when considering the full statement
+            7. If there is no surrounding context (before or after), make sure your response is a complete valid SQL statement that can be run and resolves the prompt.
+            
+            Modify the selected text now:`,
+        },
+      ],
+    })
+
+    return result.pipeDataStreamToResponse(res)
+  } catch (error) {
+    console.error('Completion error:', error)
+    return res.status(500).json({
+      error: 'Failed to generate completion',
+    })
+  }
+}

package-lock.json

@@ -13868,7 +13868,6 @@
       "version": "0.1.6",
       "resolved": "https://registry.npmjs.org/@supabase/sql-to-rest/-/sql-to-rest-0.1.6.tgz",
       "integrity": "sha512-06KgjeINtc6405XQvfnchBE1azEsU8G2NElfadmvVHKmHa5l2bFzjbtFbpaYgpgTzccHlcDmBaCgedVf2Gyl8Q==",
-      "license": "MIT",
       "dependencies": {
         "@babel/parser": "^7.24.5",
         "libpg-query": "^15.1.0",