import { PermissionAction } from '@supabase/shared-types/out/constants' import { AnimatePresence, motion } from 'framer-motion' import { last } from 'lodash' import { FileText, MessageCircleMore, Plus, WandSparkles } from 'lucide-react' import Link from 'next/link' import { useEffect, useMemo, useRef, useState } from 'react' import { toast } from 'sonner' import type { Message as MessageType } from 'ai/react' import { useChat } from 'ai/react' import { useParams } from 'common' import { Markdown } from 'components/interfaces/Markdown' import OptInToOpenAIToggle from 'components/interfaces/Organization/GeneralSettings/OptInToOpenAIToggle' import { MessageWithDebug } from 'components/interfaces/SQLEditor/AiAssistantPanel' import { DiffType } from 'components/interfaces/SQLEditor/SQLEditor.types' import { useCheckOpenAIKeyQuery } from 'data/ai/check-api-key-query' import { useSqlDebugMutation } from 'data/ai/sql-debug-mutation' import { useEntityDefinitionQuery } from 'data/database/entity-definition-query' import { useEntityDefinitionsQuery } from 'data/database/entity-definitions-query' import { useOrganizationUpdateMutation } from 'data/organizations/organization-update-mutation' import { useSendEventMutation } from 'data/telemetry/send-event-mutation' import { useCheckPermissions } from 'hooks/misc/useCheckPermissions' 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 { useFlag } from 'hooks/ui/useFlag' import { BASE_PATH, IS_PLATFORM, OPT_IN_TAGS, TELEMETRY_ACTIONS, TELEMETRY_CATEGORIES, TELEMETRY_LABELS, } from 'lib/constants' import { useAppStateSnapshot } from 'state/app-state' import { AiIconAnimation, Button, cn, DropdownMenu, DropdownMenuContent, DropdownMenuLabel, DropdownMenuRadioGroup, DropdownMenuRadioItem, DropdownMenuSeparator, DropdownMenuSub, DropdownMenuSubContent, DropdownMenuSubTrigger, DropdownMenuTrigger, SheetHeader, SheetSection, Tooltip_Shadcn_, TooltipContent_Shadcn_, TooltipTrigger_Shadcn_, } from 'ui' import { Admonition, AssistantChatForm } from 'ui-patterns' import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal' import { DocsButton } from '../DocsButton' import { ASSISTANT_SUPPORT_ENTITIES } from './AiAssistant.constants' import { SupportedAssistantEntities, SupportedAssistantQuickPromptTypes } from './AIAssistant.types' import { generatePrompt, retrieveDocsUrl } from './AIAssistant.utils' import { ContextBadge } from './ContextBadge' import { EntitiesDropdownMenu } from './EntitiesDropdownMenu' import { Message } from './Message' import { SchemasDropdownMenu } from './SchemasDropdownMenu' const ANIMATION_DURATION = 0.3 interface AIAssistantProps { id: string className?: string debugThread: MessageWithDebug[] onDiff: ({ id, diffType, sql }: { id: string; diffType: DiffType; sql: string }) => void onResetConversation: () => void } // [Joshen] For some reason I'm having issues working with dropdown menu and scroll area export const AIAssistant = ({ id, className, debugThread, onDiff, onResetConversation, }: AIAssistantProps) => { const { ref } = useParams() const project = useSelectedProject() const isOptedInToAI = useOrgOptedIntoAi() const selectedOrganization = useSelectedOrganization() const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM const disablePrompts = useFlag('disableAssistantPrompts') const { aiAssistantPanel } = useAppStateSnapshot() const { editor, entity } = aiAssistantPanel const inputRef = useRef(null) const bottomRef = useRef(null) const [value, setValue] = useState('') const [selectedDatabaseEntity, setSelectedDatabaseEntity] = useState< SupportedAssistantEntities | '' >('') const [selectedSchemas, setSelectedSchemas] = useSchemasForAi(project?.ref!) const [selectedTables, setSelectedTables] = useState<{ schema: string; name: string }[]>([]) const [contextHistory, setContextHistory] = useState<{ [key: string]: { entity: string; schemas: string[]; tables: string[] } }>({}) // [Joshen] Mainly for error handling on useChat - cause last sent messages will be voided const [assistantError, setAssistantError] = useState() const [lastSentMessage, setLastSentMessage] = useState() const [isConfirmOptInModalOpen, setIsConfirmOptInModalOpen] = useState(false) const docsUrl = retrieveDocsUrl(selectedDatabaseEntity as SupportedAssistantEntities) const entityContext = ASSISTANT_SUPPORT_ENTITIES.find((x) => x.id === selectedDatabaseEntity) const noContextAdded = selectedDatabaseEntity.length === 0 && selectedSchemas.length === 0 && selectedTables.length === 0 const { data: check } = useCheckOpenAIKeyQuery() const isApiKeySet = IS_PLATFORM || !!check?.hasKey const { data: existingDefinition } = useEntityDefinitionQuery({ id: entity?.id, type: editor, projectRef: project?.ref, connectionString: project?.connectionString, }) const { data } = useEntityDefinitionsQuery( { schemas: selectedSchemas, projectRef: project?.ref, connectionString: project?.connectionString, }, { enabled: includeSchemaMetadata } ) const tableDefinitions = selectedTables.length === 0 ? data?.map((def) => def.sql.trim()) ?? [] : data ?.filter((def) => { return selectedTables.some((table) => { return def.sql.startsWith(`CREATE TABLE ${table.schema}.${table.name}`) }) }) .map((def) => def.sql.trim()) ?? [] const entityDefinitions = includeSchemaMetadata ? tableDefinitions : undefined const { mutate: sendEvent } = useSendEventMutation() const sendTelemetryEvent = (action: string) => { sendEvent({ action, category: TELEMETRY_CATEGORIES.AI_ASSISTANT, label: TELEMETRY_LABELS.QUICK_SQL_EDITOR, }) } const { messages: chatMessages, isLoading: isChatLoading, append, } = useChat({ id, api: `${BASE_PATH}/api/ai/sql/generate-v2`, body: { entityDefinitions, context: selectedDatabaseEntity, existingSql: existingDefinition }, onError: (error) => setAssistantError(JSON.parse(error.message).error), }) const canUpdateOrganization = useCheckPermissions(PermissionAction.UPDATE, 'organizations') const { mutate: updateOrganization, isLoading: isUpdating } = useOrganizationUpdateMutation() const { isLoading: isDebugSqlLoading } = useSqlDebugMutation() const isLoading = isChatLoading || isDebugSqlLoading const messages = useMemo(() => { const merged = [ ...debugThread, ...chatMessages.map((m) => ({ ...m, isDebug: false })), ...(assistantError !== undefined && lastSentMessage !== undefined ? [lastSentMessage] : []), ] return merged.sort( (a, b) => (a.createdAt?.getTime() ?? 0) - (b.createdAt?.getTime() ?? 0) || a.role.localeCompare(b.role) ) }, [chatMessages, debugThread, assistantError, lastSentMessage]) const hasMessages = messages.length > 0 const lastMessage = last(messages) const pendingChatReply = isLoading && lastMessage?.role === 'user' const pendingDebugReply = (lastMessage as MessageWithDebug)?.isDebug && lastMessage?.role === 'assistant' && lastMessage.content === 'Thinking...' const pendingReply = pendingChatReply || pendingDebugReply const sendMessageToAssistant = (content: string) => { const payload = { role: 'user', createdAt: new Date(), content } as MessageType append(payload) setAssistantError(undefined) setLastSentMessage(payload) sendTelemetryEvent(TELEMETRY_ACTIONS.PROMPT_SUBMITTED) } const toggleSchema = (schema: string) => { if (selectedSchemas.includes(schema)) { setSelectedSchemas(selectedSchemas.filter((s) => s !== schema)) } else { const newSelectedSchemas = [...selectedSchemas, schema].sort((a, b) => a.localeCompare(b)) setSelectedSchemas(newSelectedSchemas) sendTelemetryEvent(TELEMETRY_ACTIONS.SCHEMA_CONTEXT_ADDED) } } const toggleEntity = ({ schema, name }: { schema: string; name: string }) => { const isExisting = selectedTables.find((x) => x.schema === schema && x.name === name) if (isExisting) { setSelectedTables(selectedTables.filter((x) => !(x.schema === schema && x.name === name))) } else { const newselectedTables = [...selectedTables, { schema, name }].sort( (a, b) => a.schema.localeCompare(b.schema) || a.name.localeCompare(b.name) ) setSelectedTables(newselectedTables) sendTelemetryEvent(TELEMETRY_ACTIONS.TABLE_CONTEXT_ADDED) } } const onClickQuickPrompt = (type: SupportedAssistantQuickPromptTypes) => { const prompt = generatePrompt({ type, context: selectedDatabaseEntity as any, schemas: selectedSchemas, tables: selectedTables, }) if (prompt) { setValue(prompt) sendMessageToAssistant(prompt) sendTelemetryEvent(TELEMETRY_ACTIONS.QUICK_PROMPT_SELECTED(type)) } } const confirmOptInToShareSchemaData = async () => { if (!canUpdateOrganization) { return toast.error('You do not have the required permissions to update this organization') } if (!selectedOrganization?.slug) return console.error('Organization slug is required') const existingOptInTags = selectedOrganization?.opt_in_tags ?? [] const updatedOptInTags = existingOptInTags.includes(OPT_IN_TAGS.AI_SQL) ? existingOptInTags : [...existingOptInTags, OPT_IN_TAGS.AI_SQL] updateOrganization( { slug: selectedOrganization?.slug, opt_in_tags: updatedOptInTags }, { onSuccess: () => { toast.success('Successfully opted-in') setIsConfirmOptInModalOpen(false) }, } ) } useEffect(() => { if (editor) { const mode = ASSISTANT_SUPPORT_ENTITIES.find((x) => x.id === editor) if (mode) setSelectedDatabaseEntity(mode.id) } }, [editor]) useEffect(() => { if (!isLoading) { setValue('') if (inputRef.current) inputRef.current.focus() } setTimeout( () => { if (bottomRef.current) bottomRef.current.scrollIntoView({ behavior: 'smooth' }) }, isLoading ? 100 : 500 ) }, [isLoading]) useEffect(() => { if (bottomRef.current) bottomRef.current.scrollIntoView({ behavior: 'smooth' }) }, [messages]) useEffect(() => { const lastMessage = last(messages) if (lastMessage?.role === 'user') { const entity = entityContext?.label ?? selectedDatabaseEntity setContextHistory({ ...contextHistory, [lastMessage.id]: { entity, schemas: selectedSchemas, tables: selectedTables.map((x) => selectedSchemas.length > 1 ? `${x.schema}.${x.name}` : x.name ), }, }) } // eslint-disable-next-line react-hooks/exhaustive-deps }, [messages.length]) return ( <>

Assistant

{hasMessages && ( )}
{hasMessages && ( {messages.map((m, index) => { const isFirstUserMessage = m.role === 'user' && messages.slice(0, index).every((msg) => msg.role !== 'user') return ( onDiff({ id: m.id, diffType, sql })} > {isFirstUserMessage && !includeSchemaMetadata && ( )} {isFirstUserMessage && includeSchemaMetadata && selectedSchemas.length === 0 && ( )} ) })} {assistantError !== undefined && ( )} {!isLoading && !pendingReply && assistantError === undefined && (

Please verify all responses as the Assistant can make mistakes

)} {pendingChatReply && ( )}
)}
{!hasMessages && (

How can I help you {!!entityContext ? ( <> {' '} with{' '} {entityContext.id === 'rls-policies' ? entityContext.label : `Database ${entityContext.label}`} ) : ( ' today' )} ?

)}
{disablePrompts && ( )} {!isApiKeySet && ( } /> )}
Add context for the assistant Improve the output quality of the assistant by giving it context about what you need help with {editor === null && (

Database Entity

Inform about what you're working with

setSelectedDatabaseEntity(value as SupportedAssistantEntities) } > {ASSISTANT_SUPPORT_ENTITIES.map((x) => ( {x.label} ))}
)} {includeSchemaMetadata && ( <>

Schemas

Share table definitions in the selected schemas

{selectedSchemas.length > 0 && (

Tables

Select specific tables to share definitions for

)} )}
{!!entityContext && ( setSelectedDatabaseEntity('') : undefined} /> )} {selectedSchemas.length > 0 && ( 2 ? ` and ${selectedSchemas.length - 2} other${selectedSchemas.length > 3 ? 's' : ''}` : ''}`} onRemove={() => { setSelectedSchemas([]) setSelectedTables([]) }} tooltip={ selectedSchemas.length > 2 ? ( <>

{selectedSchemas.length} schemas selected:

    {selectedSchemas.map((x) => (
  • {x}
  • ))}
) : undefined } /> )} {selectedTables.length > 0 && ( x.name) .join( ', ' )}${selectedTables.length > 2 ? ` and ${selectedTables.length - 2} other${selectedTables.length > 3 ? 's' : ''}` : ''}`} onRemove={() => setSelectedTables([])} tooltip={ selectedTables.length > 2 ? ( <>

{selectedTables.length} tables selected:

    {selectedTables.map((x) => (
  • {x.schema}.{x.name}
  • ))}
) : undefined } /> )}
textarea]:rounded-none [&>textarea]:border-0 [&>textarea]:!outline-none [&>textarea]:!ring-offset-0 [&>textarea]:!ring-0' )} loading={isLoading} disabled={!isApiKeySet || disablePrompts || isLoading} placeholder={ hasMessages ? 'Reply to the assistant...' : 'How can we help you today?' } value={value} onValueChange={(e) => setValue(e.target.value)} onSubmit={(event) => { event.preventDefault() sendMessageToAssistant(value) }} /> {!hasMessages && IS_PLATFORM && (
The Assistant is in Alpha and your prompts might be rate limited
)}
{!hasMessages && (
Suggest some{' '} {entityContext?.id === 'rls-policies' ? entityContext.label : `database ${entityContext?.label.toLowerCase()}`} Provide some examples of{' '} {entityContext?.id === 'rls-policies' ? entityContext.label : `database ${entityContext?.label.toLowerCase()}`} What are{' '} {entityContext?.id === 'rls-policies' ? entityContext.label : `database ${entityContext?.label.toLowerCase()}`} ? {docsUrl !== undefined && }
)}
setIsConfirmOptInModalOpen(false)} onConfirm={confirmOptInToShareSchemaData} loading={isUpdating} >

By opting into sending anonymous data, Supabase AI can improve the answers it shows you. This is an organization-wide setting, and affects all projects in your organization.

) }