From 3df4fe96e35dfd6b1ca8273f88f673b9c581ae32 Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Wed, 13 Nov 2024 11:18:19 +1000 Subject: [PATCH] fixes and refinements --- .../Functions/FunctionsList/FunctionList.tsx | 6 +- .../Functions/FunctionsList/FunctionsList.tsx | 3 +- .../interfaces/SQLEditor/SQLEditor.tsx | 51 +- .../ui/AIAssistantPanel/AIAssistant.tsx | 884 +++++------------- .../ui/AIAssistantPanel/AIAssistantPanel.tsx | 439 +-------- .../AIAssistantPanel/CollapsibleCodeBlock.tsx | 61 ++ .../ui/AIAssistantPanel/Message.tsx | 53 +- .../ui/AIAssistantPanel/SqlSnippet.tsx | 141 ++- apps/studio/pages/api/ai/sql/generate-v3.ts | 8 +- apps/studio/pages/api/ai/sql/tools.ts | 13 +- .../pages/project/[ref]/auth/policies.tsx | 8 +- apps/studio/state/app-state.ts | 7 + 12 files changed, 481 insertions(+), 1193 deletions(-) create mode 100644 apps/studio/components/ui/AIAssistantPanel/CollapsibleCodeBlock.tsx diff --git a/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx b/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx index b5f5bc8021d..c54b0e20994 100644 --- a/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx +++ b/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionList.tsx @@ -125,7 +125,11 @@ const FunctionList = ({ { - setAiAssistantPanel({ open: true, editor: 'functions', entity: x }) + setAiAssistantPanel({ + open: true, + initialInput: 'Help me update this function to do...', + sqlSnippets: [x.complete_statement], + }) }} > diff --git a/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionsList.tsx b/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionsList.tsx index 56b09f32faf..5652281ab82 100644 --- a/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionsList.tsx +++ b/apps/studio/components/interfaces/Database/Functions/FunctionsList/FunctionsList.tsx @@ -155,8 +155,7 @@ const FunctionsList = ({ onClick={() => setAiAssistantPanel({ open: true, - editor: 'functions', - entity: undefined, + initialInput: `Help me create a new function for the schema ${selectedSchema} that does ...`, }) } tooltip={{ diff --git a/apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx b/apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx index 1d254d90e2f..255423df1f9 100644 --- a/apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx +++ b/apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx @@ -181,7 +181,7 @@ const SQLEditor = () => { const editor = editorRef.current const monaco = monacoRef.current - const startLineNumber = hasSelection ? editor?.getSelection()?.startLineNumber ?? 0 : 0 + const startLineNumber = hasSelection ? (editor?.getSelection()?.startLineNumber ?? 0) : 0 const formattedError = error.formattedError ?? '' const lineError = formattedError.slice(formattedError.indexOf('LINE')) @@ -237,7 +237,7 @@ const SQLEditor = () => { const selection = editor.getSelection() const selectedValue = selection ? editor.getModel()?.getValueInRange(selection) : undefined const sql = snippet - ? (selectedValue || editorRef.current?.getValue()) ?? snippet.snippet.content.sql + ? ((selectedValue || editorRef.current?.getValue()) ?? snippet.snippet.content.sql) : selectedValue || editorRef.current?.getValue() formatQuery( { @@ -277,7 +277,7 @@ const SQLEditor = () => { const selectedValue = selection ? editor.getModel()?.getValueInRange(selection) : undefined const sql = snippet - ? (selectedValue || editorRef.current?.getValue()) ?? snippet.snippet.content.sql + ? ((selectedValue || editorRef.current?.getValue()) ?? snippet.snippet.content.sql) : selectedValue || editorRef.current?.getValue() let queryHasIssues = false @@ -406,26 +406,11 @@ const SQLEditor = () => { try { const snippet = snapV2.snippets[id] const result = snapV2.results[id]?.[0] - - const { solution, sql } = await debugSql({ - sql: snippet.snippet.content.sql.replace(sqlAiDisclaimerComment, '').trim(), - errorMessage: result.error.message, - entityDefinitions, + appSnap.setAiAssistantPanel({ + open: true, + sqlSnippets: [snippet.snippet.content.sql.replace(sqlAiDisclaimerComment, '').trim()], + initialInput: `Help me to debug the attached sql snippet which gives the following error: \n\n${result.error.message}`, }) - - const formattedSql = - sqlAiDisclaimerComment + - '\n\n' + - format(sql, { - language: 'postgresql', - keywordCase: 'lower', - }) - setDebugSolution(solution) - setSourceSqlDiff({ - original: snippet.snippet.content.sql, - modified: formattedSql, - }) - setSelectedDiffType(DiffType.Modification) } catch (error: unknown) { // [Joshen] There's a tendency for the SQL debug to chuck a lengthy error message // that's not relevant for the user - so we prettify it here by avoiding to return the @@ -719,7 +704,27 @@ const SQLEditor = () => { aiPanelRef.current?.expand()} + onClick={() => { + const state = getSqlEditorV2StateSnapshot() + const snippet = state.snippets[id] + const editor = editorRef.current + const selection = editor.getSelection() + const selectedValue = selection + ? editor.getModel()?.getValueInRange(selection) + : undefined + const sql = snippet + ? ((selectedValue || editorRef.current?.getValue()) ?? + snippet.snippet.content.sql) + : selectedValue || editorRef.current?.getValue() + + appSnap.setAiAssistantPanel({ + open: true, + sqlSnippets: sql ? [sql] : [], + initialInput: sql + ? `Help me make a change to the attached sql snippet` + : '', + }) + }} className={cn( 'group absolute z-10 rounded-lg right-[24px] top-4 transition-all duration-200 ease-out' )} diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index 6afeea669c0..3b8791c47c3 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx @@ -1,35 +1,18 @@ import { PermissionAction } from '@supabase/shared-types/out/constants' import { AnimatePresence, motion } from 'framer-motion' import { last, partition } from 'lodash' -import { - Box, - Code, - FileText, - MessageCircleMore, - Plus, - WandSparkles, - ChevronLeft, -} from 'lucide-react' -import Link from 'next/link' +import { Box, Code, FileText, MessageCircleMore, WandSparkles, ChevronLeft } from 'lucide-react' import { useEffect, useMemo, useRef, useState, memo } from 'react' import { toast } from 'sonner' - +import { useAppStateSnapshot } from 'state/app-state' 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' @@ -43,187 +26,91 @@ import { 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, - Heading, - SheetHeader, - SheetSection, - Tooltip_Shadcn_, - TooltipContent_Shadcn_, - TooltipTrigger_Shadcn_, -} from 'ui' +import { Button, cn } 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' -import { SQL_TEMPLATES } from 'components/interfaces/SQLEditor/SQLEditor.queries' const ANIMATION_DURATION = 0.3 import { Loading } from 'ui' +import CollapsibleCodeBlock from './CollapsibleCodeBlock' -const MemoizedMessage = memo( - ({ message, isFirstUserMessage, includeSchemaMetadata }) => { - return ( - - {isFirstUserMessage && !includeSchemaMetadata && ( - - - - )} - - ) - }, - (prevProps, nextProps) => { - // Custom comparison function to determine if the component should re-render - return ( - prevProps.message.id === nextProps.message.id && - prevProps.message.content === nextProps.message.content && - prevProps.contextHistory[prevProps.message.id] === - nextProps.contextHistory[nextProps.message.id] && - prevProps.includeSchemaMetadata === nextProps.includeSchemaMetadata && - prevProps.selectedSchemas.length === nextProps.selectedSchemas.length - ) - } -) +const MemoizedMessage = memo(({ message, isFirstUserMessage, includeSchemaMetadata }) => { + const [isConfirmOptInModalOpen, setIsConfirmOptInModalOpen] = useState(false) -type PanelType = 'assistant' | 'templates' | 'quickStart' | 'functions' + return ( + + {isFirstUserMessage && !includeSchemaMetadata && ( + + + + )} + + ) +}) -interface Panel { - title: string - component: React.ReactNode -} +MemoizedMessage.displayName = 'MemoizedMessage' interface AIAssistantProps { id: string className?: string - debugThread: MessageWithDebug[] - onDiff: ({ id, diffType, sql }: { id: string; diffType: DiffType; sql: string }) => void onResetConversation: () => void + initialMessages?: MessageType[] + sqlSnippets?: string[] + initialInput?: string } export const AIAssistant = ({ id, className, - debugThread, - onDiff, onResetConversation, + sqlSnippets: initialSqlSnippets = [], + initialMessages = [], + initialInput = '', }: AIAssistantProps) => { - const { ref } = useParams() const project = useSelectedProject() const isOptedInToAI = useOrgOptedIntoAi() const selectedOrganization = useSelectedOrganization() const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM - const [templates, quickStart] = partition(SQL_TEMPLATES, { type: 'template' }) - const [activePanel, setActivePanel] = useState('assistant') + const [sqlSnippets, setSqlSnippets] = useState(initialSqlSnippets) 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[] } - }>({}) + const [value, setValue] = useState(initialInput) const [assistantError, setAssistantError] = useState() const [lastSentMessage, setLastSentMessage] = useState() const [isConfirmOptInModalOpen, setIsConfirmOptInModalOpen] = useState(false) const [headers, setHeaders] = useState>({}) - 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: tables, - isLoading: isLoadingTables, - isSuccess: isSuccessTables, - isError: isErrorTables, - error: errorTables, - } = useTablesQuery({ + const { setAiAssistantPanel } = useAppStateSnapshot() + + const { data: tables, isLoading: isLoadingTables } = useTablesQuery({ projectRef: project?.ref, connectionString: project?.connectionString, schema: 'public', }) - 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({ @@ -243,29 +130,24 @@ export const AIAssistant = ({ api: `${BASE_PATH}/api/ai/sql/generate-v3`, headers: headers, maxSteps: 5, + initialMessages, body: { - entityDefinitions, - context: selectedDatabaseEntity, - existingSql: existingDefinition, projectRef: project?.ref, connectionString: project?.connectionString, }, onError: (error) => { console.log('error:', JSON.stringify(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 isLoading = isChatLoading const messages = useMemo(() => { const merged = [ - ...debugThread, - ...chatMessages.map((m) => ({ ...m, isDebug: false })), + ...chatMessages, ...(assistantError !== undefined && lastSentMessage !== undefined ? [lastSentMessage] : []), ] @@ -274,7 +156,7 @@ export const AIAssistant = ({ (a.createdAt?.getTime() ?? 0) - (b.createdAt?.getTime() ?? 0) || a.role.localeCompare(b.role) ) - }, [chatMessages, debugThread, assistantError, lastSentMessage]) + }, [chatMessages, assistantError, lastSentMessage]) const renderedMessages = useMemo( () => @@ -288,74 +170,28 @@ export const AIAssistant = ({ message={message} isFirstUserMessage={isFirstUserMessage} includeSchemaMetadata={includeSchemaMetadata} - selectedSchemas={selectedSchemas} - contextHistory={contextHistory} - onDiff={onDiff} /> ) }), - [messages, contextHistory, includeSchemaMetadata, selectedSchemas] + [messages, includeSchemaMetadata] ) 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 - console.log('sending:', headers) append(payload, { headers: { Authorization: headers.get('Authorization'), }, }) + setSqlSnippets([]) + setValue('') 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, entityId?: string) => { - const prompt = generatePrompt({ - type, - context: entityId || (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') @@ -380,260 +216,32 @@ export const AIAssistant = ({ ) } - const panels: Record = { - assistant: { - title: 'Assistant', - component: ( -
- {!tables || tables.length === 0 ? ( - <> -

Welcome to Supabase!

-

- I can help you get setup and even generate your entire database schema. Choose a - starting point or describe what you want to build. -

-
- {quickStart.map((template) => ( - - ))} -
- - ) : ( - <> -

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

-

- I can help you get setup and even generate your entire database schema. Choose a - starting point or describe what you want to build. -

- {activePanel === 'assistant' && !selectedDatabaseEntity && !hasMessages && ( -
- - -
- )} - {ASSISTANT_SUPPORT_ENTITIES.filter( - (e) => !selectedDatabaseEntity || e.id === selectedDatabaseEntity - ).map((entity) => ( -
-

- {entity.label} -

-
- - - - - - {/* */} -
-
- ))} - - )} -
- ), - }, - templates: { - title: 'Templates', - component: ( -
-
- {templates.map((template) => ( - - ))} -
-
- ), - }, - quickStart: { - title: 'Quick Start', - component: ( -
-
- {quickStart.map((template) => ( - - ))} -
-
- ), - }, - } - - useEffect(() => { - if (editor) { - const mode = ASSISTANT_SUPPORT_ENTITIES.find((x) => x.id === editor) - if (mode) setSelectedDatabaseEntity(mode.id) - } - }, [editor]) - useEffect(() => { const loadHeaders = async () => { const headerData = await constructHeaders() setHeaders(headerData) - console.log('auth:', headerData.get('Authorization')) } loadHeaders() + + return () => { + setAiAssistantPanel({ + initialInput: '', + sqlSnippets: [], + }) + } }, []) + useEffect(() => { + setValue(initialInput) + setSqlSnippets(initialSqlSnippets) + if (inputRef.current) { + inputRef.current.focus() + inputRef.current.setSelectionRange(initialInput.length, initialInput.length) + } + }, [initialInput, initialSqlSnippets]) + useEffect(() => { if (!isLoading) { - setValue('') if (inputRef.current) inputRef.current.focus() } @@ -649,27 +257,6 @@ export const AIAssistant = ({ 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 - ), - }, - }) - } - if (messages.length > 0) { - setActivePanel('assistant') - } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [messages.length]) - if (isLoadingTables) { return (
@@ -682,17 +269,7 @@ export const AIAssistant = ({ <>
- {activePanel !== 'assistant' && ( -
- {hasMessages && ( + {hasMessages ? ( +
+ {new Date(messages[0].createdAt || new Date()).toLocaleDateString('en-US', { + month: 'long', + day: 'numeric', + year: 'numeric', + hour: 'numeric', + minute: 'numeric', + })} +
{renderedMessages}
- )} - {!hasMessages && ( - - - {panels[activePanel].component} - - + ) : tables?.length > 0 ? ( +
+
+ +
+ + + +
+
+

How can I help you today?

+

+ I can help you get setup and even generate your entire database schema. Describe + what you want to build. +

+ {ASSISTANT_SUPPORT_ENTITIES.map((entity) => ( +
+

+ {entity.label} +

+
+ + + + + +
+
+ ))} +
+ ) : ( +
+

Welcome to Supabase!

+

+ This is the Supabase assistant which can help you create, debug and modify tables, + policies, functions and more. You can even use it to query your data using just your + words. It looks like we have a blank canvas though, so what are you looking to + build? +

+
)}
- -
-
- - - - - - - - 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 - } - /> - )} + {hasMessages && ( +
+
+ )} +
+ {sqlSnippets && sqlSnippets.length > 0 && ( +
+ {sqlSnippets.map((snippet, index) => ( + { + const newSnippets = [...sqlSnippets] + newSnippets.splice(index, 1) + setSqlSnippets(newSnippets) + }} + className="text-xs" + /> + ))} +
+ )} textarea]:rounded-none [&>textarea]:border-0 [&>textarea]:!outline-none [&>textarea]:!ring-offset-0 [&>textarea]:!ring-0' + 'z-20 [&>textarea]:border-1 [&>textarea]:rounded-md [&>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?'} + placeholder={ + hasMessages + ? 'Reply to the assistant...' + : sqlSnippets?.length > 0 + ? 'Ask a question or make a change...' + : 'How can we help you today?' + } value={value} onValueChange={(e) => setValue(e.target.value)} onSubmit={(event) => { event.preventDefault() - sendMessageToAssistant(value) + const sqlSnippetsString = sqlSnippets + .map((snippet) => '```sql\n' + snippet + '\n```') + .join('\n') + const valueWithSnippets = [value, sqlSnippetsString].filter(Boolean).join('\n\n') + sendMessageToAssistant(valueWithSnippets) }} />
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistantPanel.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistantPanel.tsx index 32d5750b8a4..8d951d2fdab 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistantPanel.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistantPanel.tsx @@ -1,452 +1,33 @@ -import { useQueryClient } from '@tanstack/react-query' -import { uniqBy } from 'lodash' -import { ChevronLeft, Command, CornerDownLeft, Loader2, X } from 'lucide-react' -import { useRouter } from 'next/router' -import { useCallback, useEffect, useRef, useState } from 'react' -import { toast } from 'sonner' - -import { useChat } from 'ai/react' -import { useParams } from 'common' -import { useIsDatabaseFunctionsAssistantEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext' -import { generateThreadMessage } from 'components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyEditorPanel.utils' -import { MessageWithDebug } from 'components/interfaces/SQLEditor/AiAssistantPanel' -import { sqlAiDisclaimerComment } from 'components/interfaces/SQLEditor/SQLEditor.constants' -import { DiffType, IStandaloneCodeEditor } from 'components/interfaces/SQLEditor/SQLEditor.types' -import { suffixWithLimit } from 'components/interfaces/SQLEditor/SQLEditor.utils' -import Results from 'components/interfaces/SQLEditor/UtilityPanel/Results' -import { useSqlDebugMutation } from 'data/ai/sql-debug-mutation' -import { databasePoliciesKeys } from 'data/database-policies/keys' -import { QueryResponseError, useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' -import { sqlKeys } from 'data/sql/keys' -import { useSendEventMutation } from 'data/telemetry/send-event-mutation' -import { usePrevious } from 'hooks/deprecated' -import { useLocalStorage } from 'hooks/misc/useLocalStorage' -import { useSelectedProject } from 'hooks/misc/useSelectedProject' -import { SqlEditor } from 'icons' -import { - BASE_PATH, - LOCAL_STORAGE_KEYS, - TELEMETRY_ACTIONS, - TELEMETRY_CATEGORIES, - TELEMETRY_LABELS, -} from 'lib/constants' -import { detectOS, uuidv4 } from 'lib/helpers' +import { useRef, useState } from 'react' +import { IStandaloneCodeEditor } from 'components/interfaces/SQLEditor/SQLEditor.types' +import { uuidv4 } from 'lib/helpers' import { useAppStateSnapshot } from 'state/app-state' -import { - Button, - cn, - Sheet, - SheetContent, - SheetFooter, - SheetHeader, - Tooltip_Shadcn_, - TooltipContent_Shadcn_, - TooltipTrigger_Shadcn_, -} from 'ui' -import { Admonition } from 'ui-patterns' -import CodeEditor from '../CodeEditor/CodeEditor' +import { cn } from 'ui' import { AIAssistant } from './AIAssistant' -import { generateCTA, generatePlaceholder, generateTitle, validateQuery } from './AIAssistant.utils' -import { ASSISTANT_SUPPORT_ENTITIES } from './AiAssistant.constants' -import { useEntityDefinitionQuery } from 'data/database/entity-definition-query' export const AiAssistantPanel = () => { - const os = detectOS() - const router = useRouter() - const { ref } = useParams() - const project = useSelectedProject() - const queryClient = useQueryClient() const { aiAssistantPanel, setAiAssistantPanel } = useAppStateSnapshot() - const isEnabled = useIsDatabaseFunctionsAssistantEnabled() - const { open, editor, content, entity } = aiAssistantPanel - const previousEditor = usePrevious(editor) - - const [isAcknowledged, setIsAcknowledged] = useLocalStorage( - LOCAL_STORAGE_KEYS.SQL_SCRATCH_PAD_BANNER_ACKNOWLEDGED, - false - ) + const { open, editor, content, entity, initialMessages, initialInput, sqlSnippets } = + aiAssistantPanel const [chatId, setChatId] = useState(() => uuidv4()) - const [error, setError] = useState() - const [results, setResults] = useState(undefined) - const [showResults, setShowResults] = useState(false) - const [showWarning, setShowWarning] = useState(false) - const [debugThread, setDebugThread] = useState([]) - - const { data: existingDefinition } = useEntityDefinitionQuery({ - id: entity?.id, - type: editor, - projectRef: project?.ref, - connectionString: project?.connectionString, - }) // [Joshen] JFYI I'm opting to just have the assistant always show and not toggle-able // I don't really see any negatives of keeping it open (or benefits from hiding) tbh // const [showAssistant, setShowAssistant] = useState(true) - const title = generateTitle(editor, entity) - const placeholder = generatePlaceholder(editor, entity, existingDefinition) - const ctaText = generateCTA(editor) const editorRef = useRef() - const editorModel = editorRef.current?.getModel() - - const numResults = (results ?? []).length - const [errorHeader, ...errorContent] = - (error?.formattedError?.split('\n') ?? [])?.filter((x: string) => x.length > 0) ?? [] - const entityContext = ASSISTANT_SUPPORT_ENTITIES.find((x) => x.id === editor) - - const { mutate: sendEvent } = useSendEventMutation() - const sendTelemetryEvent = (action: string) => { - sendEvent({ - action, - category: TELEMETRY_CATEGORIES.AI_ASSISTANT, - label: TELEMETRY_LABELS.QUICK_SQL_EDITOR, - }) - } - - const { mutate: executeSql, isLoading: isExecuting } = useExecuteSqlMutation({ - onSuccess: async (res) => { - // [Joshen] If in a specific editor context mode, assume that intent was to create/update - // a database entity - so close it once success. Otherwise it's in Quick SQL mode and we - // show the results. Currently though it assumes we're "creating", thinking need to support "updating" too - if (editor !== null) { - switch (editor) { - case 'functions': - await queryClient.invalidateQueries(sqlKeys.query(ref, ['functions-list'])) - break - case 'rls-policies': - await queryClient.invalidateQueries(databasePoliciesKeys.list(ref)) - break - } - - toast.success(`Successfully created ${entityContext?.name}!`) - setAiAssistantPanel({ open: false }) - } else { - setShowResults(true) - setResults(res.result) - } - }, - onError: (error) => { - setError(error) - setResults([]) - }, - }) - - const { append } = useChat({ - id: chatId, - api: `${BASE_PATH}/api/ai/sql/generate-v2`, - // [Joshen] Don't need entity definitions if calling here cause support action - // via AI here is just to explain code segment, no need for entity definitions context - body: {}, - }) - - const { mutateAsync: debugSql } = useSqlDebugMutation({ onError: () => {} }) - - const [view, setView] = useState<'assistant' | 'editor'>('assistant') - - const isEditor = view === 'editor' - - const onExecuteSql = (skipValidation = false) => { - setError(undefined) - setShowWarning(false) - - const query = editorRef.current?.getValue() ?? '' - if (query.length === 0) return - - if (editor !== undefined && !skipValidation) { - // [Joshen] Some basic validation logic - const validated = validateQuery(editor, query) - if (!validated) { - return setShowWarning(true) - } - } - - executeSql({ - sql: suffixWithLimit(query, 100), - projectRef: project?.ref, - connectionString: project?.connectionString, - handleError: (error) => { - throw error - }, - }) - } - - const onFixWithAssistant = async () => { - if (!error) return - - const messageId = uuidv4() - const query = editorRef.current?.getValue() ?? '' - - const assistantMessageBefore = generateThreadMessage({ - id: messageId, - content: 'Thinking...', - isDebug: true, - }) - - setDebugThread([...debugThread, assistantMessageBefore]) - sendTelemetryEvent(TELEMETRY_ACTIONS.FIX_WITH_ASSISTANT) - - try { - const { solution, sql } = await debugSql({ - sql: query, - errorMessage: error.message, - }) - - const assistantMessageAfter = generateThreadMessage({ - id: messageId, - content: `${solution}\n\`\`\`sql\n${sql}\n\`\`\``, - isDebug: true, - }) - const cleanedMessages = uniqBy([...debugThread, assistantMessageAfter], (m) => m.id) - setDebugThread(cleanedMessages) - } catch (error) { - const assistantMessageAfter = generateThreadMessage({ - id: messageId, - content: `Hmm, Sorry but I'm unable to find a solution for the error that you're facing`, - isDebug: true, - }) - const cleanedMessages = uniqBy([...debugThread, assistantMessageAfter], (m) => m.id) - setDebugThread(cleanedMessages) - } - } - - const onExplainSql = (value: string) => { - append({ - role: 'user', - createdAt: new Date(), - content: ` - Can you explain this section to me in more detail?\n - ${value} - `.trim(), - }) - sendTelemetryEvent(TELEMETRY_ACTIONS.EXPLAIN_CODE) - } - - const updateEditorWithCheckForDiff = useCallback( - ({ id, diffType, sql }: { id: string; diffType: DiffType; sql: string }) => { - 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: `${sqlAiDisclaimerComment}\n\n${sql}`, - range: editorModel.getFullModelRange(), - }, - ]) - } else { - const currentSql = editorRef.current?.getValue() - editorRef.current?.executeEdits('apply-ai-message', [ - { - text: `${currentSql}\n\n${sqlAiDisclaimerComment}\n\n${sql}`, - range: editorModel.getFullModelRange(), - }, - ]) - } - setView('editor') - }, - [editorModel] - ) - - const onTogglePanel = () => { - if (open) { - // [Joshen] Save the code content when closing - should allow users to continue from - // where they left off IF they are opening the assistant again in the same "editor" context - const existingValue = editorRef.current?.getValue() ?? '' - setAiAssistantPanel({ open: false, content: existingValue }) - } else { - setAiAssistantPanel({ open: true, editor: null, entity: undefined }) - } - } - - useEffect(() => { - if (open) { - // [Joshen] Only reset the assistant if the editor changed - if (previousEditor !== editor) { - setChatId(uuidv4()) - setError(undefined) - setShowWarning(false) - setDebugThread([]) - setResults(undefined) - setShowResults(false) - } else { - // [Joshen] Repopulate the code editor with the content from where the user left off - // setTimeout is just to give time for the editorRef to get associated - setTimeout(() => { - const editorModel = editorRef.current?.getModel() - if (content && editorModel) { - editorRef.current?.executeEdits('apply-ai-message', [ - { text: content, range: editorModel.getFullModelRange() }, - ]) - } - }, 100) - } - } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [open]) - - // [Joshen] Just to test the concept of a universal assistant of sorts - useEffect(() => { - const handler = (e: KeyboardEvent) => { - if (e.metaKey && e.code === 'KeyI') onTogglePanel() - } - if (project !== undefined && isEnabled) window.addEventListener('keydown', handler) - return () => window.removeEventListener('keydown', handler) - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [project, isEnabled, open]) - - // [Joshen] Whenever the deps change recalculate the height of the editor - useEffect(() => { - editorRef.current?.layout({ width: 0, height: 0 }) - window.requestAnimationFrame(() => editorRef.current?.layout()) - }, [error, showWarning, isExecuting, numResults, showResults]) return !open ? null : (
setChatId(uuidv4())} /> -
-
-
- - {editor === null && !isAcknowledged && ( - - - Queries written here will not be saved and results are limited up to only 100 rows - - - - } - /> - )} - {showWarning && ( - -

Are you sure you want to run this query?

- - - } - /> - )} - {results !== undefined && results.length > 0 && ( - <> -
- -
-
-

- {results.length} rows - {results.length >= 100 && ` (Limited to only 100 rows)`} -

- -
- - )} - {results !== undefined && results.length === 0 && ( -
-

Success. No rows returned.

-
- )} -
- - -
-
-
-
) } diff --git a/apps/studio/components/ui/AIAssistantPanel/CollapsibleCodeBlock.tsx b/apps/studio/components/ui/AIAssistantPanel/CollapsibleCodeBlock.tsx new file mode 100644 index 00000000000..0a60ef1055f --- /dev/null +++ b/apps/studio/components/ui/AIAssistantPanel/CollapsibleCodeBlock.tsx @@ -0,0 +1,61 @@ +import { ChevronDown, ChevronUp, X } from 'lucide-react' +import { useState } from 'react' +import { Button, CodeBlock, CodeBlockProps, cn } from 'ui' + +interface CollapsibleCodeBlockProps extends CodeBlockProps { + onRemove?: () => void +} + +const CollapsibleCodeBlock = ({ onRemove, ...props }: CollapsibleCodeBlockProps) => { + const [isExpanded, setIsExpanded] = useState(false) + + const codeString = (props.value || props.children) as string + const firstLine = isExpanded + ? codeString + : codeString?.substring(0, codeString.indexOf('\n')) || codeString + + return ( +
+
+
+
+ ) +} + +export default CollapsibleCodeBlock diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx index 4c601a1bfb4..523612572fa 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.tsx @@ -9,18 +9,18 @@ import { SqlSnippet } from './SqlSnippet' import { DiffType } from 'components/interfaces/SQLEditor/SQLEditor.types' import { useProfile } from 'lib/profile' import { MessagePre } from './MessagePre' +import { motion } from 'framer-motion' +import CollapsibleCodeBlock from './CollapsibleCodeBlock' interface MessageProps { name?: string role: 'function' | 'system' | 'user' | 'assistant' | 'data' | 'tool' content?: string createdAt?: number - isDebug?: boolean isSelected?: boolean isLoading?: boolean + readOnly?: boolean action?: React.ReactNode - context?: { entity: string; schemas: string[]; tables: string[] } - onDiff?: (type: DiffType, s: string) => void variant?: 'default' | 'warning' } @@ -29,43 +29,35 @@ export const Message = function Message({ role, content, createdAt, - isDebug, isLoading, + readOnly, isSelected = false, - context, children, action = null, variant = 'default', - onDiff = noop, }: PropsWithChildren) { const { profile } = useProfile() const isUser = role === 'user' - const formattedContext = - context !== undefined - ? Object.entries(context) - .filter(([_, value]) => value.length > 0) - .map(([key, value]) => { - return `${key.charAt(0).toUpperCase() + key.slice(1)}: ${Array.isArray(value) ? value.join(', ') : value}` - }) - .join(' • ') - : undefined - if (!content) return null return ( -
+ {children}
{variant === 'warning' && } - {isDebug && Debug request} {action} { - return + return readOnly ? ( +
+ +
+ ) : ( + + ) }, ol: (props: any) => { return
    {props.children}
@@ -92,10 +98,7 @@ export const Message = function Message({ > {content} - {/* {role === 'user' && context !== undefined && ( - {formattedContext} - )} */}
-
+ ) } diff --git a/apps/studio/components/ui/AIAssistantPanel/SqlSnippet.tsx b/apps/studio/components/ui/AIAssistantPanel/SqlSnippet.tsx index af168af8e4c..9ff54ab694b 100644 --- a/apps/studio/components/ui/AIAssistantPanel/SqlSnippet.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/SqlSnippet.tsx @@ -1,5 +1,5 @@ import { useState, useEffect, useCallback } from 'react' -import { Code, Play, Edit, DatabaseIcon } from 'lucide-react' +import { Code, Play, Edit, DatabaseIcon, FileWarning } from 'lucide-react' import { Button, CodeBlock, cn } from 'ui' import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' import { useSelectedProject } from 'hooks/misc/useSelectedProject' @@ -14,12 +14,14 @@ import useNewQuery from 'components/interfaces/SQLEditor/hooks' interface SqlSnippetWrapperProps { sql: string isLoading?: boolean + readOnly?: boolean } interface ParsedSqlProps { sql: string props: any title: string + readOnly?: boolean } const isReadOnlySelect = (query: string): boolean => { @@ -65,7 +67,7 @@ const isReadOnlySelect = (query: string): boolean => { }) } -const SqlSnippetWrapper = ({ sql, isLoading }: SqlSnippetWrapperProps) => { +const SqlSnippetWrapper = ({ sql, isLoading, readOnly = false }: SqlSnippetWrapperProps) => { let formatted = (sql || [''])[0] const propsMatch = formatted.match(/--\s*props:\s*(\{[^}]+\})/) const props = propsMatch ? JSON.parse(propsMatch[1]) : {} @@ -96,7 +98,7 @@ const SqlSnippetWrapper = ({ sql, isLoading }: SqlSnippetWrapperProps) => {
{ xAxis={props.xAxis} yAxis={props.yAxis} title={title} + readOnly={readOnly} />
) } -export const SqlCard = ({ sql, isChart, xAxis, yAxis, title }: ParsedSqlProps) => { - const [showCode, setShowCode] = useState(!isReadOnlySelect(sql)) +export const SqlCard = ({ + sql, + isChart, + xAxis, + yAxis, + title, + readOnly = false, +}: ParsedSqlProps) => { + const [showCode, setShowCode] = useState(readOnly || !isReadOnlySelect(sql)) const [showResults, setShowResults] = useState(false) const [results, setResults] = useState() const project = useSelectedProject() const [error, setError] = useState() + const [showWarning, setShowWarning] = useState(false) const { newQuery } = useNewQuery() const { mutate: executeSql, isLoading: isExecuting } = useExecuteSqlMutation({ @@ -122,17 +133,25 @@ export const SqlCard = ({ sql, isChart, xAxis, yAxis, title }: ParsedSqlProps) = console.log('SQL executed successfully:', res) setShowResults(true) setResults(res.result) + setShowWarning(false) }, onError: (error) => { setError(error) setResults([]) + setShowWarning(false) }, }) const handleExecute = useCallback(() => { - if (!project?.ref || !sql) return - setError(undefined) + if (!project?.ref || !sql || readOnly) return + if (!isReadOnlySelect(sql)) { + setShowCode(true) + setShowWarning(true) + return + } + + setError(undefined) executeSql({ sql: suffixWithLimit(sql, 100), projectRef: project?.ref, @@ -142,55 +161,97 @@ export const SqlCard = ({ sql, isChart, xAxis, yAxis, title }: ParsedSqlProps) = setResults([]) }, }) - }, [project?.ref, project?.connectionString, sql, executeSql]) + }, [project?.ref, project?.connectionString, sql, executeSql, readOnly]) useEffect(() => { - if (isReadOnlySelect(sql) && !results) { + if (isReadOnlySelect(sql) && !results && !readOnly) { handleExecute() } - }, [sql, handleExecute]) + }, [sql, handleExecute, readOnly]) const handleEdit = () => { - newQuery(sql, title) + if (!readOnly) { + newQuery(sql, title) + } } const [errorHeader, ...errorContent] = (error?.formattedError?.split('\n') ?? [])?.filter((x: string) => x.length > 0) ?? [] - console.log('results:', results) - return (
- -

{title}

+ {showWarning ? ( +
+ +

+ This query contains write operations. Are you sure you want to execute it? +

+
+ + +
+
+ ) : ( + <> + +

{title}

-
- + {!readOnly && ( +
+ - + - -
+ +
+ )} + + )}
{showCode && ( @@ -199,8 +260,6 @@ export const SqlCard = ({ sql, isChart, xAxis, yAxis, title }: ParsedSqlProps) = language="sql" className={cn( 'max-h-96 block !bg-transparent !py-3 !px-3.5 prose dark:prose-dark border-0 border-t text-foreground !rounded-none w-full', - // change the look of the code block. The flex hack is so that the code is wrapping since - // every word is a separate span '[&>code]:m-0 [&>code>span]:flex [&>code>span]:flex-wrap [&>code]:block [&>code>span]:text-foreground' )} hideLineNumbers diff --git a/apps/studio/pages/api/ai/sql/generate-v3.ts b/apps/studio/pages/api/ai/sql/generate-v3.ts index a9b8acf86a3..2115fb6eb2a 100644 --- a/apps/studio/pages/api/ai/sql/generate-v3.ts +++ b/apps/studio/pages/api/ai/sql/generate-v3.ts @@ -59,7 +59,11 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse) # You write row level security policies. - Your purpose is to generate a policy with the constraints given by the user. First retrieve the schema or schemas that are most relevant. Then retrieve RLS policy examples. + 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. # 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 @@ -72,7 +76,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse) 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. - Here are the existing database schemas: ${schemas} + Here are the existing database schema names you can retrieve: ${schemas} `, messages, tools: getTools(projectRef, connectionString, authorization), diff --git a/apps/studio/pages/api/ai/sql/tools.ts b/apps/studio/pages/api/ai/sql/tools.ts index 30cd0279204..59bd7d56404 100644 --- a/apps/studio/pages/api/ai/sql/tools.ts +++ b/apps/studio/pages/api/ai/sql/tools.ts @@ -4,12 +4,12 @@ import { executeSql } from 'data/sql/execute-sql-query' import { getEntityDefinitionsQuery } from 'data/database/entity-definitions-query' import { getDatabasePolicies } from 'data/database-policies/database-policies-query' import { stripIndent } from 'common-tags' -type DatabasePoliciesData = components['schemas']['PostgresPolicy'] +import { DatabasePoliciesData } from 'data/database-policies/database-policies-query' export const getTools = (projectRef: string, connectionString: string, authorization?: string) => { return { getSchema: tool({ - description: 'Get the schema definition for a given schema', + description: 'Get more information about one or more schemas', parameters: z.object({ schemas: z.array(z.string()).describe('The schema names to get the definitions for'), }), @@ -38,8 +38,9 @@ export const getTools = (projectRef: string, connectionString: string, authoriza } }, }), - getRlsExamples: tool({ - description: 'Get examples and instructions on how to write RLS policies', + getRlsKnowledge: tool({ + description: + 'Get existing policies and examples and instructions on how to write RLS policies', parameters: z.object({ schemas: z.array(z.string()).describe('The schema names to get the policies for'), }), @@ -55,7 +56,6 @@ export const getTools = (projectRef: string, connectionString: string, authoriza ...(authorization && { Authorization: authorization }), } ) - console.log('rls:', data) const formattedPolicies = data .map( (policy: DatabasePoliciesData) => ` @@ -71,8 +71,7 @@ export const getTools = (projectRef: string, connectionString: string, authoriza return stripIndent` You're a Supabase Postgres expert in writing row level security policies. Your purpose is to - generate a policy with the constraints given by the user. You will be provided a schema - on which the policy should be applied. + generate a policy with the constraints given by the user. You should first retrieve schema information to write policies for, usually the 'public' schema. The output should use the following instructions: - The generated SQL must be valid SQL. diff --git a/apps/studio/pages/project/[ref]/auth/policies.tsx b/apps/studio/pages/project/[ref]/auth/policies.tsx index e603d3734b9..3cba1b29e97 100644 --- a/apps/studio/pages/project/[ref]/auth/policies.tsx +++ b/apps/studio/pages/project/[ref]/auth/policies.tsx @@ -21,6 +21,7 @@ import { EXCLUDED_SCHEMAS } from 'lib/constants/schemas' import type { NextPageWithLayout } from 'types' import { Button, Input } from 'ui' import { DocsButton } from 'components/ui/DocsButton' +import { useAppStateSnapshot } from 'state/app-state' /** * Filter tables by table name and policy name @@ -64,6 +65,7 @@ const AuthPoliciesPage: NextPageWithLayout = () => { }>() const { schema = 'public', search: searchString = '' } = params const { project } = useProjectContext() + const { setAiAssistantPanel, aiAssistantPanel } = useAppStateSnapshot() const [selectedTable, setSelectedTable] = useState() const [showPolicyAiEditor, setShowPolicyAiEditor] = useState(false) @@ -148,8 +150,10 @@ const AuthPoliciesPage: NextPageWithLayout = () => { hasTables={tables.length > 0} isLocked={isLocked} onSelectCreatePolicy={(table: string) => { - setShowPolicyAiEditor(true) - setSelectedTable(table) + setAiAssistantPanel({ + open: true, + initialInput: `Help me create an RLS policy for the table ${table} on the ${schema} schema`, + }) }} onSelectEditPolicy={(policy) => { setSelectedPolicyToEdit(policy) diff --git a/apps/studio/state/app-state.ts b/apps/studio/state/app-state.ts index 71f10820099..453bf2eda0e 100644 --- a/apps/studio/state/app-state.ts +++ b/apps/studio/state/app-state.ts @@ -3,6 +3,7 @@ import { proxy, snapshot, useSnapshot } from 'valtio' import { SupportedAssistantEntities } from 'components/ui/AIAssistantPanel/AIAssistant.types' import { LOCAL_STORAGE_KEYS } from 'lib/constants' import { LOCAL_STORAGE_KEYS as COMMON_LOCAL_STORAGE_KEYS } from 'common' +import type { Message as MessageType } from 'ai/react' const EMPTY_DASHBOARD_HISTORY: { sql?: string @@ -21,6 +22,9 @@ export type CommonDatabaseEntity = { type AiAssistantPanelType = { open: boolean + initialMessages?: MessageType[] + initialInput?: string + sqlSnippets?: string[] editor?: SupportedAssistantEntities | null // Raw string content for the monaco editor, currently used to retain where the user left off when toggling off the panel content?: string @@ -109,6 +113,9 @@ export const appState = proxy({ aiAssistantPanel: { open: false, + initialMessages: [], + sqlSnippets: [], + initialInput: '', editor: null, content: '', entity: undefined,