diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index 3eb400a6ecb..26c2a81aef3 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx @@ -8,11 +8,14 @@ import { toast } from 'sonner' import type { Message as MessageType } from 'ai/react' import { useChat } from 'ai/react' +import { useParams } from 'common' +import { subscriptionHasHipaaAddon } from 'components/interfaces/Billing/Subscription/Subscription.utils' import OptInToOpenAIToggle from 'components/interfaces/Organization/GeneralSettings/OptInToOpenAIToggle' import { SQL_TEMPLATES } from 'components/interfaces/SQLEditor/SQLEditor.queries' import { useCheckOpenAIKeyQuery } from 'data/ai/check-api-key-query' import { constructHeaders } from 'data/fetchers' import { useOrganizationUpdateMutation } from 'data/organizations/organization-update-mutation' +import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query' import { useTablesQuery } from 'data/tables/tables-query' import { useSendEventMutation } from 'data/telemetry/send-event-mutation' import { useCheckPermissions } from 'hooks/misc/useCheckPermissions' @@ -30,6 +33,7 @@ import { } from 'lib/constants' import uuidv4 from 'lib/uuid' import { useAppStateSnapshot } from 'state/app-state' +import { useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2' import { AiIconAnimation, Button, @@ -44,8 +48,6 @@ import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal' import AIOnboarding from './AIOnboarding' import CollapsibleCodeBlock from './CollapsibleCodeBlock' import { Message } from './Message' -import { getSqlEditorV2StateSnapshot, useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2' -import { useParams } from 'common' const MemoizedMessage = memo(({ message }: { message: MessageType }) => { return ( @@ -72,6 +74,7 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP const router = useRouter() const project = useSelectedProject() const { id: snippetId } = useParams() + const isOptedInToAI = useOrgOptedIntoAi() const selectedOrganization = useSelectedOrganization() const includeSchemaMetadata = isOptedInToAI || !IS_PLATFORM @@ -98,6 +101,9 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP const snippet = snippets[snippetId ?? ''] const snippetContent = snippet?.snippet?.content?.sql + const { data: subscription } = useOrgSubscriptionQuery({ orgSlug: selectedOrganization?.slug }) + const hasHipaaAddon = subscriptionHasHipaaAddon(subscription) + const { data: tables, isLoading: isLoadingTables } = useTablesQuery({ projectRef: project?.ref, connectionString: project?.connectionString, @@ -122,8 +128,10 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP id, api: `${BASE_PATH}/api/ai/sql/generate-v3`, maxSteps: 5, + // [Joshen] Not currently used atm, but initialMessages will be for... initialMessages: initialMessages as unknown as MessageType[], body: { + includeSchemaMetadata, projectRef: project?.ref, connectionString: project?.connectionString, }, @@ -297,16 +305,22 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP - + {!hasHipaaAddon && ( + + )} )} @@ -468,11 +482,13 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP )} + {showFade && (
-
+
)} +
{sqlSnippets && sqlSnippets.length > 0 && (
@@ -510,14 +526,19 @@ export const AIAssistant = ({ id, className, onResetConversation }: AIAssistantP onValueChange={(e) => setValue(e.target.value)} onSubmit={(event) => { event.preventDefault() - const sqlSnippetsString = - sqlSnippets?.map((snippet) => '```sql\n' + snippet + '\n```').join('\n') || '' - const valueWithSnippets = [value, sqlSnippetsString].filter(Boolean).join('\n\n') - sendMessageToAssistant(valueWithSnippets) + if (includeSchemaMetadata) { + const sqlSnippetsString = + sqlSnippets?.map((snippet) => '```sql\n' + snippet + '\n```').join('\n') || '' + const valueWithSnippets = [value, sqlSnippetsString].filter(Boolean).join('\n\n') + sendMessageToAssistant(valueWithSnippets) + } else { + sendMessageToAssistant(value) + } }} />
+ { +import { getDatabasePolicies } from 'data/database-policies/database-policies-query' +import { getEntityDefinitionsQuery } from 'data/database/entity-definitions-query' +import { executeSql } from 'data/sql/execute-sql-query' + +export const getTools = ({ + projectRef, + connectionString, + authorization, + includeSchemaMetadata, +}: { + projectRef: string + connectionString: string + authorization?: string + includeSchemaMetadata: boolean +}) => { return { getSchema: tool({ description: 'Get more information about one or more schemas', @@ -14,22 +24,24 @@ export const getTools = (projectRef: string, connectionString: string, authoriza schemas: z.array(z.string()).describe('The schema names to get the definitions for'), }), execute: async ({ schemas }) => { - // retrieve schema console.log('schema:', schemas) try { - const result = await executeSql( - { - projectRef, - connectionString, - sql: getEntityDefinitionsQuery({ schemas }), - }, - undefined, - { - 'Content-Type': 'application/json', - ...(authorization && { Authorization: authorization }), - } - ) + const result = includeSchemaMetadata + ? await executeSql( + { + projectRef, + connectionString, + sql: getEntityDefinitionsQuery({ schemas }), + }, + undefined, + { + 'Content-Type': 'application/json', + ...(authorization && { Authorization: authorization }), + } + ) + : { result: [] } + console.log('response2:', result) return result } catch (error) { @@ -45,20 +57,23 @@ export const getTools = (projectRef: string, connectionString: string, authoriza schemas: z.array(z.string()).describe('The schema names to get the policies for'), }), execute: async ({ schemas }) => { - const data = await getDatabasePolicies( - { - projectRef, - connectionString, - }, - undefined, - { - 'Content-Type': 'application/json', - ...(authorization && { Authorization: authorization }), - } - ) + const data = includeSchemaMetadata + ? await getDatabasePolicies( + { + projectRef, + connectionString, + }, + undefined, + { + 'Content-Type': 'application/json', + ...(authorization && { Authorization: authorization }), + } + ) + : [] + const formattedPolicies = data .map( - (policy: DatabasePoliciesData) => ` + (policy) => ` Policy Name: "${policy.name}" Action: ${policy.action} Roles: ${policy.roles.join(', ')} @@ -306,7 +321,7 @@ export const getTools = (projectRef: string, connectionString: string, authoriza This prevents the policy \`( (select auth.uid()) = user_id )\` from running for any \`anon\` users, since the execution stops at the \`to authenticated\` step. - Here are my existing policies: ${formattedPolicies} + ${data.length > 0 ? `Here are my existing policies: ${formattedPolicies}` : ''} ` }, }),