konnosDashboard
useAIOptInForm.ts111 lines · main
1import { zodResolver } from '@hookform/resolvers/zod'
2import { PermissionAction } from '@supabase/shared-types/out/constants'
3import { useQueryClient } from '@tanstack/react-query'
4import { LOCAL_STORAGE_KEYS } from 'common'
5import { useForm } from 'react-hook-form'
6import { toast } from 'sonner'
7import * as z from 'zod'
8
9import { useOrganizationUpdateMutation } from '@/data/organizations/organization-update-mutation'
10import { invalidateOrganizationsQuery } from '@/data/organizations/organizations-query'
11import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions'
12import { useLocalStorageQuery } from '@/hooks/misc/useLocalStorage'
13import { getAiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
14import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization'
15import { OPT_IN_TAGS } from '@/lib/constants'
16import type { ResponseError } from '@/types'
17
18// Shared schema definition
19export const AIOptInSchema = z.object({
20 aiOptInLevel: z.enum(['disabled', 'schema', 'schema_and_log', 'schema_and_log_and_data'], {
21 required_error: 'AI Opt-in level selection is required',
22 }),
23})
24
25export type AIOptInFormValues = z.infer<typeof AIOptInSchema>
26
27/**
28 * Hook to manage the AI Opt-In form state and submission logic.
29 * Optionally takes an onSuccess callback (e.g., to close a modal).
30 */
31export const useAIOptInForm = (onSuccessCallback?: () => void) => {
32 const queryClient = useQueryClient()
33 const { data: selectedOrganization } = useSelectedOrganizationQuery()
34 const { can: canUpdateOrganization } = useAsyncCheckPermissions(
35 PermissionAction.UPDATE,
36 'organizations'
37 )
38
39 const [, setUpdatedOptInSinceMCP] = useLocalStorageQuery(
40 LOCAL_STORAGE_KEYS.AI_ASSISTANT_MCP_OPT_IN,
41 false
42 )
43
44 const { mutate: updateOrganization, isPending: isUpdating } = useOrganizationUpdateMutation()
45
46 const form = useForm<AIOptInFormValues>({
47 resolver: zodResolver(AIOptInSchema as any),
48 defaultValues: {
49 aiOptInLevel: getAiOptInLevel(selectedOrganization?.opt_in_tags),
50 },
51 })
52
53 const onSubmit = async (values: AIOptInFormValues) => {
54 if (!canUpdateOrganization) {
55 return toast.error('You do not have the required permissions to update this organization')
56 }
57 if (!selectedOrganization?.slug) {
58 console.error('Organization slug is required')
59 return toast.error('Failed to update settings: Organization not found.')
60 }
61 const existingOptInTags = selectedOrganization?.opt_in_tags ?? []
62
63 let updatedOptInTags = existingOptInTags.filter(
64 (tag: string) =>
65 tag !== OPT_IN_TAGS.AI_SQL &&
66 tag !== (OPT_IN_TAGS.AI_DATA ?? 'AI_DATA') &&
67 tag !== (OPT_IN_TAGS.AI_LOG ?? 'AI_LOG')
68 )
69
70 if (
71 values.aiOptInLevel === 'schema' ||
72 values.aiOptInLevel === 'schema_and_log' ||
73 values.aiOptInLevel === 'schema_and_log_and_data'
74 ) {
75 updatedOptInTags.push(OPT_IN_TAGS.AI_SQL)
76 }
77 if (
78 values.aiOptInLevel === 'schema_and_log' ||
79 values.aiOptInLevel === 'schema_and_log_and_data'
80 ) {
81 updatedOptInTags.push(OPT_IN_TAGS.AI_LOG)
82 }
83 if (values.aiOptInLevel === 'schema_and_log_and_data') {
84 updatedOptInTags.push(OPT_IN_TAGS.AI_DATA)
85 }
86
87 updatedOptInTags = [...new Set(updatedOptInTags)]
88
89 updateOrganization(
90 { slug: selectedOrganization.slug, opt_in_tags: updatedOptInTags },
91 {
92 onSuccess: () => {
93 invalidateOrganizationsQuery(queryClient)
94 toast.success('Successfully updated AI opt-in settings')
95 setUpdatedOptInSinceMCP(true)
96 onSuccessCallback?.() // Call optional callback on success
97 },
98 onError: (error: ResponseError) => {
99 toast.error(`Failed to update settings: ${error.message}`)
100 },
101 }
102 )
103 }
104
105 return {
106 form,
107 onSubmit,
108 isUpdating,
109 currentOptInLevel: getAiOptInLevel(selectedOrganization?.opt_in_tags),
110 }
111}