konnosDashboard
AIAssistant.tsx604 lines · main
1// @ts-nocheck
2import type { UIMessage as MessageType } from '@ai-sdk/react'
3import { useChat } from '@ai-sdk/react'
4import { lastAssistantMessageIsCompleteWithApprovalResponses } from 'ai'
5import { LOCAL_STORAGE_KEYS, useFlag } from 'common'
6import { useParams, useSearchParamsShallow } from 'common/hooks'
7import { AnimatePresence, motion } from 'framer-motion'
8import { Eraser, Pencil, X } from 'lucide-react'
9import { useRouter } from 'next/router'
10import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
11import { Button, cn, KeyboardShortcut } from 'ui'
12import { Admonition } from 'ui-patterns'
13
14import AlertError from '../AlertError'
15import { ButtonTooltip } from '../ButtonTooltip'
16import { ErrorBoundary } from '../ErrorBoundary/ErrorBoundary'
17import { ASSISTANT_ERRORS } from './AiAssistant.constants'
18import type { SqlSnippet } from './AIAssistant.types'
19import {
20 hasPendingToolApproval,
21 onErrorChat,
22 resolvePendingToolApprovalsAsDenied,
23} from './AIAssistant.utils'
24import { AIAssistantHeader } from './AIAssistantHeader'
25import { AIOnboarding } from './AIOnboarding'
26import { AssistantChatForm } from './AssistantChatForm'
27import {
28 Conversation,
29 ConversationContent,
30 ConversationScrollButton,
31} from './elements/Conversation'
32import { Message } from './Message'
33import { Markdown } from '@/components/interfaces/Markdown'
34import { SIDEBAR_KEYS } from '@/components/layouts/ProjectLayout/LayoutSidebar/LayoutSidebarProvider'
35import { useCheckOpenAIKeyQuery } from '@/data/ai/check-api-key-query'
36import { useRateMessageMutation } from '@/data/ai/rate-message-mutation'
37import { useTablesQuery } from '@/data/tables/tables-query'
38import { useCheckEntitlements } from '@/hooks/misc/useCheckEntitlements'
39import { useLocalStorageQuery } from '@/hooks/misc/useLocalStorage'
40import { useOrgAiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
41import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization'
42import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
43import {
44 DEFAULT_ASSISTANT_BASE_MODEL_ID,
45 defaultAssistantModelId,
46 isAssistantBaseModelId,
47 isKnownAssistantModelId,
48} from '@/lib/ai/model.utils'
49import { IS_PLATFORM } from '@/lib/constants'
50import { uuidv4 } from '@/lib/helpers'
51import { useTrack } from '@/lib/telemetry/track'
52import type { AssistantModel } from '@/state/ai-assistant-state'
53import { useAiAssistantState, useAiAssistantStateSnapshot } from '@/state/ai-assistant-state'
54import { SHORTCUT_IDS } from '@/state/shortcuts/registry'
55import { useShortcut } from '@/state/shortcuts/useShortcut'
56import { useSidebarManagerSnapshot } from '@/state/sidebar-manager-state'
57import { useSqlEditorV2StateSnapshot } from '@/state/sql-editor-v2'
58
59interface AIAssistantProps {
60 initialMessages?: MessageType[] | undefined
61 className?: string
62}
63
64export const AIAssistant = ({ className }: AIAssistantProps) => {
65 const router = useRouter()
66 const { id: entityId } = useParams()
67 const { data: project } = useSelectedProjectQuery()
68 const searchParams = useSearchParamsShallow()
69
70 const { data: selectedOrganization, isPending: isLoadingOrganization } =
71 useSelectedOrganizationQuery()
72
73 useShortcut(SHORTCUT_IDS.AI_ASSISTANT_CANCEL_EDIT, () => cancelEdit())
74
75 const disablePrompts = useFlag('disableAssistantPrompts')
76 const { snippets } = useSqlEditorV2StateSnapshot()
77 const snap = useAiAssistantStateSnapshot()
78 const state = useAiAssistantState()
79 const { activeSidebar, closeSidebar } = useSidebarManagerSnapshot()
80
81 const { hasAccess: hasAccessToAdvanceModel, isLoading: isLoadingEntitlements } =
82 useCheckEntitlements('assistant.advance_model')
83
84 const selectedModel = useMemo<AssistantModel>(() => {
85 // While entitlements are loading, use the stored model without enforcing access
86 if (isLoadingEntitlements) {
87 return snap.model ?? DEFAULT_ASSISTANT_BASE_MODEL_ID
88 }
89
90 const defaultModel = defaultAssistantModelId(hasAccessToAdvanceModel)
91 const model = snap.model ?? defaultModel
92
93 if (!isKnownAssistantModelId(model)) return defaultModel
94 if (!hasAccessToAdvanceModel && !isAssistantBaseModelId(model)) {
95 return DEFAULT_ASSISTANT_BASE_MODEL_ID
96 }
97
98 return model
99 }, [isLoadingEntitlements, hasAccessToAdvanceModel, snap.model])
100
101 const [updatedOptInSinceMCP] = useLocalStorageQuery(
102 LOCAL_STORAGE_KEYS.AI_ASSISTANT_MCP_OPT_IN,
103 false
104 )
105
106 const inputRef = useRef<HTMLTextAreaElement>(null)
107
108 const { aiOptInLevel, isHipaaProjectDisallowed } = useOrgAiOptInLevel()
109 const showMetadataWarning =
110 IS_PLATFORM &&
111 !!selectedOrganization &&
112 (aiOptInLevel === 'disabled' || aiOptInLevel === 'schema')
113
114 // Add a ref to store the last user message
115 const lastUserMessageRef = useRef<MessageType | null>(null)
116
117 // Keep latest selected organization to avoid stale values in useChat transport
118 const selectedOrganizationRef = useRef(selectedOrganization)
119 useEffect(() => {
120 selectedOrganizationRef.current = selectedOrganization
121 }, [selectedOrganization])
122
123 const [value, setValue] = useState<string>(snap.initialInput || '')
124 const [editingMessageId, setEditingMessageId] = useState<string | null>(null)
125 const [isResubmitting, setIsResubmitting] = useState(false)
126 const [messageRatings, setMessageRatings] = useState<Record<string, 'positive' | 'negative'>>({})
127
128 const { data: check, isSuccess } = useCheckOpenAIKeyQuery()
129 const isApiKeySet = !!check?.hasKey
130
131 const { mutateAsync: rateMessage } = useRateMessageMutation()
132
133 const isInSQLEditor = router.pathname.includes('/sql/[id]')
134 const snippet = snippets[entityId ?? '']
135 const snippetContent = snippet?.snippet?.content?.unchecked_sql
136
137 const { data: tables } = useTablesQuery(
138 {
139 projectRef: project?.ref,
140 connectionString: project?.connectionString,
141 schema: 'public',
142 },
143 { enabled: isApiKeySet }
144 )
145
146 const currentTable = tables?.find((t) => t.id.toString() === entityId)
147 const currentSchema = searchParams?.get('schema') ?? 'public'
148
149 // Update context in state
150 useEffect(() => {
151 state.setContext({
152 projectRef: project?.ref,
153 orgSlug: selectedOrganizationRef.current?.slug,
154 connectionString: project?.connectionString ?? '',
155 })
156 }, [project?.ref, project?.connectionString, selectedOrganizationRef.current?.slug, state])
157
158 const track = useTrack()
159
160 const {
161 messages: chatMessages,
162 status: chatStatus,
163 error,
164 sendMessage,
165 setMessages,
166 addToolApprovalResponse,
167 stop,
168 regenerate,
169 } = useChat({
170 id: snap.activeChatId,
171 ...(snap.activeChatId && snap.chatInstances[snap.activeChatId]
172 ? { chat: snap.chatInstances[snap.activeChatId] }
173 : {}),
174 sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,
175 onError: onErrorChat,
176 })
177
178 const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming'
179 const hasPendingApproval = hasPendingToolApproval(chatMessages)
180 const isChatInputDisabled = !isApiKeySet || disablePrompts || isLoadingOrganization
181
182 const deleteMessageFromHere = useCallback(
183 (messageId: string) => {
184 // Find the message index in current chatMessages
185 const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId)
186 if (messageIndex === -1) return
187
188 if (isChatLoading) stop()
189
190 snap.deleteMessagesAfter(messageId, { includeSelf: true })
191
192 const updatedMessages = chatMessages.slice(0, messageIndex)
193 setMessages(updatedMessages)
194 },
195 [snap, setMessages, chatMessages, isChatLoading, stop]
196 )
197
198 const editMessage = useCallback(
199 (messageId: string) => {
200 const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId)
201 if (messageIndex === -1) return
202
203 // Target message
204 const messageToEdit = chatMessages[messageIndex]
205
206 // Activate editing mode
207 setEditingMessageId(messageId)
208 const textContent =
209 messageToEdit.parts
210 ?.filter((part) => part.type === 'text')
211 .map((part) => part.text)
212 .join('') ?? ''
213 setValue(textContent)
214
215 setTimeout(() => {
216 if (inputRef.current) {
217 inputRef?.current?.focus()
218
219 // [Joshen] This is just to make the cursor go to the end of the text when focusing
220 const val = inputRef.current.value
221 inputRef.current.value = ''
222 inputRef.current.value = val
223 }
224 }, 100)
225 },
226 [chatMessages, setValue]
227 )
228
229 const cancelEdit = useCallback(() => {
230 setEditingMessageId(null)
231 setValue('')
232 }, [setValue])
233
234 const handleRateMessage = useCallback(
235 async (messageId: string, rating: 'positive' | 'negative', reason?: string) => {
236 if (!project?.ref || !selectedOrganization?.slug) return
237
238 // Optimistically update UI
239 setMessageRatings((prev) => ({ ...prev, [messageId]: rating }))
240
241 try {
242 const result = await rateMessage({
243 rating,
244 messages: chatMessages,
245 messageId,
246 projectRef: project.ref,
247 orgSlug: selectedOrganization.slug,
248 reason,
249 spanId: state.messageSpanIds[messageId],
250 })
251
252 track('assistant_message_rating_submitted', {
253 rating,
254 category: result.category,
255 ...(reason && { reason }),
256 chatId: state.activeChatId,
257 })
258 } catch (error) {
259 console.error('Failed to rate message:', error)
260 // Rollback on error
261 setMessageRatings((prev) => {
262 const { [messageId]: _, ...rest } = prev
263 return rest
264 })
265 }
266 },
267 [chatMessages, project?.ref, selectedOrganization?.slug, rateMessage, track, state]
268 )
269
270 const isContextExceededError =
271 error &&
272 (error.message?.includes('context_length_exceeded') ||
273 error.message?.includes('exceeds the context window'))
274
275 const renderedMessages = useMemo(
276 () =>
277 chatMessages.map((message, index) => {
278 const isBeingEdited = editingMessageId === message.id
279 const isAfterEditedMessage = editingMessageId
280 ? chatMessages.findIndex((m) => m.id === editingMessageId) < index
281 : false
282 const isLastMessage = index === chatMessages.length - 1
283
284 return (
285 <Message
286 id={message.id}
287 key={message.id}
288 message={message}
289 isLoading={chatStatus === 'submitted' || chatStatus === 'streaming'}
290 readOnly={message.role === 'user'}
291 addToolApprovalResponse={addToolApprovalResponse}
292 onDelete={deleteMessageFromHere}
293 onEdit={editMessage}
294 isAfterEditedMessage={isAfterEditedMessage}
295 isBeingEdited={isBeingEdited}
296 onCancelEdit={cancelEdit}
297 isLastMessage={isLastMessage}
298 onRate={handleRateMessage}
299 rating={messageRatings[message.id] ?? null}
300 />
301 )
302 }),
303 [
304 chatMessages,
305 deleteMessageFromHere,
306 editMessage,
307 cancelEdit,
308 editingMessageId,
309 chatStatus,
310 addToolApprovalResponse,
311 handleRateMessage,
312 messageRatings,
313 ]
314 )
315
316 const hasMessages = chatMessages.length > 0
317
318 const sendMessageToAssistant = (finalContent: string) => {
319 if (editingMessageId) {
320 // Handling when the user is in edit mode
321 // delete the message(s) from the chat just like the delete button
322 setIsResubmitting(true)
323 deleteMessageFromHere(editingMessageId)
324 setEditingMessageId(null)
325 }
326
327 const payload = {
328 role: 'user',
329 createdAt: new Date(),
330 parts: [{ type: 'text', text: finalContent }],
331 id: uuidv4(),
332 } as MessageType
333
334 snap.clearSqlSnippets()
335 lastUserMessageRef.current = payload
336 if (hasPendingApproval && !editingMessageId) {
337 setMessages(resolvePendingToolApprovalsAsDenied(chatMessages))
338 }
339 sendMessage(payload, {
340 body: {
341 schema: currentSchema,
342 table: currentTable?.name,
343 },
344 })
345 setValue('')
346
347 if (finalContent.includes('Help me to debug')) {
348 track('assistant_debug_submitted', { chatId: snap.activeChatId })
349 } else {
350 track('assistant_prompt_submitted', { chatId: snap.activeChatId })
351 }
352 }
353
354 const handleClearMessages = () => {
355 if (isChatLoading) stop()
356 snap.clearMessages()
357 setMessages([])
358 lastUserMessageRef.current = null
359 setEditingMessageId(null)
360 }
361
362 useEffect(() => {
363 // Keep "Thinking" visible while stopping and resubmitting during edit
364 // Only clear once the new response actually starts streaming (or errors)
365 if (isResubmitting && (chatStatus === 'streaming' || !!error)) {
366 setIsResubmitting(false)
367 }
368 }, [isResubmitting, chatStatus, error])
369
370 useEffect(() => {
371 setValue(snap.initialInput || '')
372 if (inputRef.current && snap.initialInput) {
373 inputRef.current.focus()
374 inputRef.current.setSelectionRange(snap.initialInput.length, snap.initialInput.length)
375 }
376 }, [snap.initialInput])
377
378 useEffect(() => {
379 const isOpen = activeSidebar?.id === SIDEBAR_KEYS.AI_ASSISTANT
380 if (isOpen && isInSQLEditor && !!snippetContent) {
381 snap.setSqlSnippets([{ label: 'Current Query', content: snippetContent }])
382 }
383 // eslint-disable-next-line react-hooks/exhaustive-deps
384 }, [activeSidebar?.id, isInSQLEditor, snippetContent])
385
386 return (
387 <ErrorBoundary
388 message="Something went wrong with the AI Assistant"
389 sentryContext={{
390 component: 'AIAssistant',
391 feature: 'AI Assistant Panel',
392 projectRef: project?.ref,
393 organizationSlug: selectedOrganization?.slug,
394 }}
395 actions={[
396 {
397 label: 'Clear messages and refresh',
398 onClick: () => {
399 handleClearMessages()
400 window.location.reload()
401 },
402 },
403 ]}
404 >
405 <div className={cn('flex flex-col h-full w-full md:h-full max-h-dvh', className)}>
406 <AIAssistantHeader
407 isChatLoading={isChatLoading}
408 onNewChat={snap.newChat}
409 onCloseAssistant={() => closeSidebar(SIDEBAR_KEYS.AI_ASSISTANT)}
410 showMetadataWarning={showMetadataWarning}
411 updatedOptInSinceMCP={updatedOptInSinceMCP}
412 isHipaaProjectDisallowed={isHipaaProjectDisallowed}
413 aiOptInLevel={aiOptInLevel}
414 />
415 {hasMessages ? (
416 <Conversation className={cn('flex-1')}>
417 <ConversationContent className="w-full px-7 py-8 mb-10">
418 {renderedMessages}
419 {error && (
420 <>
421 <AlertError
422 error={
423 isContextExceededError
424 ? ASSISTANT_ERRORS['context-exceeded']
425 : IS_PLATFORM
426 ? ASSISTANT_ERRORS['default']
427 : error
428 }
429 showErrorPrefix={false}
430 showInstructions={false}
431 subject="Sorry, I'm having trouble responding right now."
432 additionalActions={
433 <div className="flex items-center gap-x-2 mr-auto">
434 {isContextExceededError ? (
435 <Button
436 type="default"
437 size="tiny"
438 onClick={() => snap.newChat()}
439 className="text-xs"
440 >
441 New chat
442 </Button>
443 ) : (
444 <>
445 <Button
446 type="default"
447 size="tiny"
448 onClick={() => regenerate()}
449 className="text-xs"
450 >
451 Retry
452 </Button>
453 <ButtonTooltip
454 type="default"
455 size="tiny"
456 onClick={handleClearMessages}
457 className="w-7 h-7"
458 icon={<Eraser />}
459 tooltip={{ content: { side: 'bottom', text: 'Clear messages' } }}
460 />
461 </>
462 )}
463 </div>
464 }
465 />
466 </>
467 )}
468 {isChatLoading && (
469 <motion.span
470 animate={{ opacity: [1, 0] }}
471 transition={{ duration: 1, repeat: Infinity, ease: 'linear' }}
472 className="inline-block w-1.5 h-4 bg-foreground-lighter mt-4"
473 />
474 )}
475 <p className="text-center text-xs text-foreground-muted mt-6">
476 Briven AI may not always produce correct answers. Double check responses.
477 </p>
478 </ConversationContent>
479 <ConversationScrollButton />
480 </Conversation>
481 ) : (
482 <AIOnboarding
483 sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
484 suggestions={
485 snap.suggestions as
486 | { title?: string; prompts?: { label: string; description: string }[] }
487 | undefined
488 }
489 onValueChange={(val) => setValue(val)}
490 onFocusInput={() => inputRef.current?.focus()}
491 />
492 )}
493
494 <AnimatePresence>
495 {editingMessageId && (
496 <motion.div
497 initial={{ opacity: 0 }}
498 animate={{ opacity: 1 }}
499 exit={{ opacity: 0 }}
500 className="pointer-events-none z-10 -mt-24"
501 >
502 <div className="h-24 w-full bg-linear-to-t from-background to-transparent relative">
503 <motion.div
504 className="absolute left-1/2 z-20 bottom-8 pointer-events-auto"
505 variants={{
506 hidden: { y: 5, opacity: 0 },
507 show: { y: 0, opacity: 1 },
508 }}
509 transition={{ duration: 0.1 }}
510 initial="hidden"
511 animate="show"
512 exit="hidden"
513 >
514 <div className="-translate-x-1/2 bg-alternative dark:bg-muted border rounded-md px-3 py-2 min-w-[180px] flex items-center justify-between gap-x-2">
515 <div className="flex items-center gap-x-2 text-sm text-foreground">
516 <Pencil size={14} />
517 <span>Editing message</span>
518 </div>
519 <ButtonTooltip
520 type="outline"
521 size="tiny"
522 icon={<X size={14} />}
523 onClick={cancelEdit}
524 className="w-6 h-6 p-0"
525 title="Cancel editing"
526 aria-label="Cancel editing"
527 tooltip={{
528 content: { side: 'top', text: <KeyboardShortcut keys={['Meta', 'Esc']} /> },
529 }}
530 />
531 </div>
532 </motion.div>
533 </div>
534 </motion.div>
535 )}
536 </AnimatePresence>
537
538 <div className="px-3 pb-3 z-20 relative">
539 {disablePrompts && (
540 <Admonition
541 showIcon={false}
542 type="default"
543 title="Assistant has been temporarily disabled"
544 description="We're currently looking into getting it back online"
545 />
546 )}
547
548 {isSuccess && !isApiKeySet && (
549 <Admonition
550 type="default"
551 title="OpenAI API key not set"
552 description={
553 <Markdown
554 content={
555 'Add your `OPENAI_API_KEY` to your environment variables to use the AI Assistant.'
556 }
557 />
558 }
559 />
560 )}
561
562 <AssistantChatForm
563 textAreaRef={inputRef}
564 className={cn(
565 'z-20 [&>form>textarea]:text-base [&>form>textarea]:md:text-sm [&>form>textarea]:border [&>form>textarea]:rounded-md [&>form>textarea]:outline-hidden! [&>form>textarea]:ring-offset-0! [&>form>textarea]:ring-0!'
566 )}
567 loading={isChatLoading}
568 isEditing={!!editingMessageId}
569 disabled={isChatInputDisabled}
570 placeholder={
571 hasMessages
572 ? 'Ask a follow up question...'
573 : (snap.sqlSnippets ?? [])?.length > 0
574 ? 'Ask a question or make a change...'
575 : 'Chat to Postgres...'
576 }
577 value={value}
578 onValueChange={(e) => setValue(e.target.value)}
579 onSubmit={(finalMessage) => {
580 sendMessageToAssistant(finalMessage)
581 }}
582 onStop={() => {
583 stop()
584 // to save partial responses from the AI
585 const lastMessage = chatMessages[chatMessages.length - 1]
586 if (lastMessage && lastMessage.role === 'assistant') {
587 state.updateMessage(lastMessage)
588 }
589 }}
590 sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
591 onRemoveSnippet={(index) => {
592 const newSnippets = [...(snap.sqlSnippets ?? [])]
593 newSnippets.splice(index, 1)
594 snap.setSqlSnippets(newSnippets)
595 }}
596 includeSnippetsInMessage={aiOptInLevel !== 'disabled'}
597 selectedModel={selectedModel}
598 onSelectModel={(model) => snap.setModel(model)}
599 />
600 </div>
601 </div>
602 </ErrorBoundary>
603 )
604}