konnosDashboard
AIAssistant.utils.ts158 lines · main
1import { isToolUIPart, type UIMessage } from 'ai'
2import { toast } from 'sonner'
3
4import { SAFE_FUNCTIONS } from './AiAssistant.constants'
5import { authKeys } from '@/data/auth/keys'
6import { databaseExtensionsKeys } from '@/data/database-extensions/keys'
7import { databaseIndexesKeys } from '@/data/database-indexes/keys'
8import { databasePoliciesKeys } from '@/data/database-policies/keys'
9import { databaseTriggerKeys } from '@/data/database-triggers/keys'
10import { databaseKeys } from '@/data/database/keys'
11import { enumeratedTypesKeys } from '@/data/enumerated-types/keys'
12import { handleError } from '@/data/fetchers'
13import { tableKeys } from '@/data/tables/keys'
14import { tryParseJson } from '@/lib/helpers'
15import { ResponseError } from '@/types'
16
17export type MutationCategory = 'functions' | 'rls-policies'
18
19// [Joshen] This is just very basic identification, but possible can extend perhaps
20export const identifyQueryType = (query: string): MutationCategory | undefined => {
21 const formattedQuery = query.toLowerCase().replaceAll('\n', ' ')
22 if (
23 formattedQuery.includes('create function') ||
24 formattedQuery.includes('create or replace function')
25 ) {
26 return 'functions'
27 } else if (formattedQuery.includes('create policy') || formattedQuery.includes('alter policy')) {
28 return 'rls-policies'
29 }
30 return undefined
31}
32
33// Check for function calls that aren't in the safe list
34/** @deprecated [Joshen] Ideally we move away from this as this isn't a scalable way to deduce */
35export const containsUnknownFunction = (query: string) => {
36 const normalizedQuery = query.trim().toLowerCase()
37 const functionCallRegex = /\w+\s*\(/g
38 const functionCalls = normalizedQuery.match(functionCallRegex) || []
39
40 return functionCalls.some((func) => {
41 const isReadOnlyFunc = SAFE_FUNCTIONS.some((safeFunc) => func.trim().toLowerCase() === safeFunc)
42 return !isReadOnlyFunc
43 })
44}
45
46/** @deprecated
47 * [Joshen] This isn't really a scalable way to reduce this behaviour, we now have support
48 * for a readonly connection string which we can use this to run queries, and is a much
49 * clearer way to deduce if the query is read only or not
50 */
51export const isReadOnlySelect = (query: string): boolean => {
52 const normalizedQuery = query.trim().toLowerCase()
53
54 // Check if it starts with SELECT
55 if (!normalizedQuery.startsWith('select')) return false
56
57 // List of keywords that indicate write operations
58 const writeOperations = ['insert', 'update', 'delete', 'alter', 'drop', 'create', 'replace']
59
60 // Words that may appear in column names etc
61 const allowedPatterns = ['created', 'inserted', 'updated', 'deleted', 'truncate']
62
63 // Check for any write operations
64 const hasWriteOperation = writeOperations.some((op) => {
65 // Ignore if part of allowed pattern
66 const isAllowed = allowedPatterns.some(
67 (allowed) => normalizedQuery.includes(allowed) && allowed.includes(op)
68 )
69 return !isAllowed && normalizedQuery.includes(op)
70 })
71 if (hasWriteOperation) return false
72
73 const hasUnknownFunction = containsUnknownFunction(normalizedQuery)
74 if (hasUnknownFunction) return false
75
76 return true
77}
78
79export const hasPendingToolApproval = (messages: Pick<UIMessage, 'role' | 'parts'>[]) => {
80 return messages.some((message) => {
81 if (message.role !== 'assistant') return false
82
83 return message.parts?.some((part) => isToolUIPart(part) && part.state === 'approval-requested')
84 })
85}
86
87export const resolvePendingToolApprovalsAsDenied = (messages: UIMessage[]): UIMessage[] => {
88 return messages.map((message) => {
89 if (message.role !== 'assistant') return message
90
91 const parts = message.parts?.map((part) => {
92 if (!isToolUIPart(part) || part.state !== 'approval-requested') return part
93
94 return {
95 ...part,
96 state: 'output-denied',
97 approval: {
98 id: part.approval.id,
99 approved: false,
100 reason: 'Skipped because the user sent a follow-up message.',
101 },
102 } as UIMessage['parts'][number]
103 })
104
105 return { ...message, parts } as UIMessage
106 })
107}
108
109const getContextKey = (pathname: string) => {
110 const [, , , ...rest] = pathname.split('/')
111 const key = rest.join('/')
112 return key
113}
114
115export const getContextualInvalidationKeys = ({
116 ref,
117 pathname,
118 schema = 'public',
119}: {
120 ref: string
121 pathname: string
122 schema?: string
123}) => {
124 const key = getContextKey(pathname)
125
126 return (
127 (
128 {
129 'auth/users': [authKeys.usersInfinite(ref)],
130 'auth/policies': [databasePoliciesKeys.list(ref)],
131 'database/functions': [databaseKeys.databaseFunctions(ref)],
132 'database/tables': [tableKeys.list(ref, schema, true), tableKeys.list(ref, schema, false)],
133 'database/triggers': [databaseTriggerKeys.list(ref)],
134 'database/types': [enumeratedTypesKeys.list(ref)],
135 'database/extensions': [databaseExtensionsKeys.list(ref)],
136 'database/indexes': [databaseIndexesKeys.list(ref, schema)],
137 } as const
138 )[key] ?? []
139 )
140}
141
142export const onErrorChat = (error: Error) => {
143 const parsedError = error ? tryParseJson(error.message) : undefined
144
145 try {
146 handleError(parsedError?.error || parsedError || error)
147 } catch (e: any) {
148 if (e instanceof ResponseError) {
149 toast.error(e.message)
150 } else if (e instanceof Error) {
151 toast.error(e.message)
152 } else if (typeof e === 'string') {
153 toast.error(e)
154 } else {
155 toast.error('An unknown error occurred')
156 }
157 }
158}