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}` : ''}
`
},
}),