Prompt and tool refactoring (supabase#37500)

* try a really long context window to maximize caching

* update examples

* attempt to update packages and useChat

* update endpoints

* update zod

* zod

* update to v5

* message update

* Revert "zod"

This reverts commit ec39bac.

* revert zod

* zod i

* fix complete endpoints

* remove async

* change to content

* type cleanup

* Revert the package bumps to rebuild them.

* Bump zod to 2.25.76 in all packages.

* Bump openai in all packages.

* Bump ai and ai-related packages.

* Remove unneeded files.

* Fix the rest of the migration stuff.

* Prettier fixes.

* add policy list tool

* refactor

* ai sdk 5 fixes

* refactor complete endpoint

* edge function prompt

* remove example

* slight prompt change

* Minor clean up

* More clean up

---------

Co-authored-by: Jordi Enric <[email protected]>
Co-authored-by: Ivan Vasilov <[email protected]>
Co-authored-by: Joshen Lim <[email protected]>

Author: 4 people

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

@@ -1,8 +1,7 @@
+import { useCompletion } from '@ai-sdk/react'
 import type { Monaco } from '@monaco-editor/react'
 import { useQueryClient } from '@tanstack/react-query'
-import { useCompletion } from '@ai-sdk/react'
-import { AnimatePresence, motion } from 'framer-motion'
-import { ChevronUp, Command, Loader2 } from 'lucide-react'
+import { ChevronUp, Loader2 } from 'lucide-react'
 import dynamic from 'next/dynamic'
 import { useRouter } from 'next/router'
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
@@ -459,11 +458,12 @@ export const SQLEditor = () => {
     completion,
     isLoading: isCompletionLoading,
   } = useCompletion({
-    api: `${BASE_PATH}/api/ai/sql/complete-v2`,
+    api: `${BASE_PATH}/api/ai/code/complete`,
     body: {
       projectRef: project?.ref,
       connectionString: project?.connectionString,
-      includeSchemaMetadata,
+      language: 'sql',
+      orgSlug: org?.slug,
     },
     onError: (error) => {
       toast.error(`Failed to generate SQL: ${error.message}`)

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

@@ -1,4 +1,4 @@
-import { Check, ChevronDown, Edit, MessageSquare, Plus, Trash, X } from 'lucide-react'
+import { Check, ChevronDown, Edit, Plus, Trash, X } from 'lucide-react'
 import { useState } from 'react'
 
 import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state'
@@ -18,17 +18,12 @@ import {
   PopoverTrigger_Shadcn_,
   ScrollArea,
 } from 'ui'
-import { ButtonTooltip } from '../ButtonTooltip'
 
 interface AIAssistantChatSelectorProps {
-  className?: string
   disabled?: boolean
 }
 
-export const AIAssistantChatSelector = ({
-  className,
-  disabled = false,
-}: AIAssistantChatSelectorProps) => {
+export const AIAssistantChatSelector = ({ disabled = false }: AIAssistantChatSelectorProps) => {
   const snap = useAiAssistantStateSnapshot()
   const currentChat = snap.activeChat?.name
 
@@ -211,7 +206,7 @@ export const AIAssistantChatSelector = ({
                 disabled={disabled}
               >
                 <Plus size={14} strokeWidth={1.5} />
-                <span>New chat</span>
+                <span>Start a new chat</span>
               </CommandItem_Shadcn_>
             </CommandGroup_Shadcn_>
           </CommandList_Shadcn_>

apps/studio/components/ui/AIEditor/index.tsx

@@ -19,7 +19,7 @@ interface AIEditorProps {
   aiMetadata?: {
     projectRef?: string
     connectionString?: string | null
-    includeSchemaMetadata?: boolean
+    orgSlug?: string
   }
   initialPrompt?: string
   readOnly?: boolean

apps/studio/components/ui/EditorPanel/EditorPanel.tsx

@@ -12,7 +12,7 @@ import Results from 'components/interfaces/SQLEditor/UtilityPanel/Results'
 import { SqlRunButton } from 'components/interfaces/SQLEditor/UtilityPanel/RunButton'
 import { useSqlTitleGenerateMutation } from 'data/ai/sql-title-mutation'
 import { QueryResponseError, useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
-import { useOrgAiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
+import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization'
 import { useSelectedProjectQuery } from 'hooks/misc/useSelectedProject'
 import { BASE_PATH } from 'lib/constants'
 import { uuidv4 } from 'lib/helpers'
@@ -93,7 +93,7 @@ export const EditorPanel = ({
   const { profile } = useProfile()
   const snapV2 = useSqlEditorV2StateSnapshot()
   const { mutateAsync: generateSqlTitle } = useSqlTitleGenerateMutation()
-  const { includeSchemaMetadata } = useOrgAiOptInLevel()
+  const { data: org } = useSelectedOrganizationQuery()
 
   const [isSaving, setIsSaving] = useState(false)
   const [error, setError] = useState<QueryResponseError>()
@@ -328,11 +328,11 @@ export const EditorPanel = ({
               language="pgsql"
               value={currentValue}
               onChange={handleChange}
-              aiEndpoint={`${BASE_PATH}/api/ai/sql/complete-v2`}
+              aiEndpoint={`${BASE_PATH}/api/ai/code/complete`}
               aiMetadata={{
                 projectRef: project?.ref,
                 connectionString: project?.connectionString,
-                includeSchemaMetadata,
+                orgSlug: org?.slug,
               }}
               initialPrompt={initialPrompt}
               options={{

apps/studio/components/ui/FileExplorerAndEditor/FileExplorerAndEditor.tsx

@@ -28,7 +28,7 @@ interface FileExplorerAndEditorProps {
   aiMetadata?: {
     projectRef?: string
     connectionString?: string | null
-    includeSchemaMetadata?: boolean
+    orgSlug?: string
   }
 }
 

apps/studio/lib/ai/org-ai-details.ts

@@ -0,0 +1,46 @@
+import { getOrganizations } from 'data/organizations/organizations-query'
+import { getProjects } from 'data/projects/projects-query'
+import { getAiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
+
+export const getOrgAIDetails = async ({
+  orgSlug,
+  authorization,
+  projectRef,
+}: {
+  orgSlug: string
+  authorization: string
+  projectRef: string
+}) => {
+  const [organizations, projects] = await Promise.all([
+    getOrganizations({
+      headers: {
+        'Content-Type': 'application/json',
+        ...(authorization && { Authorization: authorization }),
+      },
+    }),
+    getProjects({
+      headers: {
+        'Content-Type': 'application/json',
+        ...(authorization && { Authorization: authorization }),
+      },
+    }),
+  ])
+
+  const selectedOrg = organizations.find((org) => org.slug === orgSlug)
+  const selectedProject = projects.find(
+    (project) => project.ref === projectRef || project.preview_branch_refs.includes(projectRef)
+  )
+
+  // If the project is not in the organization specific by the org slug, return an error
+  if (selectedProject?.organization_slug !== selectedOrg?.slug) {
+    throw new Error('Project and organization do not match')
+  }
+
+  const aiOptInLevel = getAiOptInLevel(selectedOrg?.opt_in_tags)
+  const isLimited = selectedOrg?.plan.id === 'free'
+
+  return {
+    aiOptInLevel,
+    isLimited,
+  }
+}