konnosDashboard
ai-assistant-state.tsx659 lines · main
1// @ts-nocheck
2import { Chat, type UIMessage as MessageType } from '@ai-sdk/react'
3import { DefaultChatTransport, lastAssistantMessageIsCompleteWithApprovalResponses } from 'ai'
4import { LOCAL_STORAGE_KEYS } from 'common'
5import { DBSchema, IDBPDatabase, openDB } from 'idb'
6import { debounce } from 'lodash'
7import { createContext, PropsWithChildren, useContext, useEffect, useState } from 'react'
8import { v4 as uuidv4 } from 'uuid'
9import { proxy, ref, snapshot, subscribe, useSnapshot } from 'valtio'
10
11import { constructHeaders } from '@/data/fetchers'
12import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
13import { prepareMessagesForAPI } from '@/lib/ai/message-utils'
14import { isKnownAssistantModelId } from '@/lib/ai/model.utils'
15import type { AssistantModelId } from '@/lib/ai/model.utils'
16import { BASE_PATH, IS_PLATFORM } from '@/lib/constants'
17
18type SuggestionsType = {
19 title: string
20 prompts?: { label: string; description: string }[]
21}
22
23export type AssistantMessageType = MessageType
24
25export type SqlSnippet = string | { label: string; content: string }
26
27export type AssistantModel = AssistantModelId
28
29type ChatSession = {
30 id: string
31 name: string
32 messages: AssistantMessageType[]
33 createdAt: Date
34 updatedAt: Date
35}
36
37export type AiAssistantContext = {
38 projectRef?: string
39 orgSlug?: string
40 connectionString?: string
41}
42
43type AiAssistantData = {
44 initialInput: string
45 sqlSnippets?: SqlSnippet[]
46 suggestions?: SuggestionsType
47 tables: { schema: string; name: string }[]
48 chats: Record<string, ChatSession>
49 activeChatId?: string
50 model?: AssistantModel
51 context: AiAssistantContext
52}
53
54// Data structure stored in IndexedDB
55type StoredAiAssistantState = {
56 projectRef: string
57 activeChatId?: string
58 chats: Record<string, ChatSession>
59 model?: AssistantModel
60}
61
62const INITIAL_AI_ASSISTANT: AiAssistantData = {
63 initialInput: '',
64 sqlSnippets: undefined,
65 suggestions: undefined,
66 tables: [],
67 chats: {},
68 activeChatId: undefined,
69 model: undefined,
70 context: {},
71}
72
73const DB_NAME = 'ai-assistant-db'
74const DB_VERSION = 1
75const STORE_NAME = 'assistantState'
76
77interface AiAssistantDB extends DBSchema {
78 [STORE_NAME]: {
79 key: string
80 value: StoredAiAssistantState
81 }
82}
83
84async function openAiDb(): Promise<IDBPDatabase<AiAssistantDB>> {
85 return openDB<AiAssistantDB>(DB_NAME, DB_VERSION, {
86 upgrade(db) {
87 if (!db.objectStoreNames.contains(STORE_NAME)) {
88 db.createObjectStore(STORE_NAME, { keyPath: 'projectRef' })
89 }
90 },
91 })
92}
93
94async function getAiState(projectRef: string): Promise<StoredAiAssistantState | undefined> {
95 if (!projectRef) return undefined
96 try {
97 const db = await openAiDb()
98 return await db.get(STORE_NAME, projectRef)
99 } catch (error) {
100 console.error('Failed to get AI state from IndexedDB:', error)
101 return undefined
102 }
103}
104
105async function saveAiState(state: StoredAiAssistantState): Promise<void> {
106 if (!state.projectRef) return
107 try {
108 const db = await openAiDb()
109 await db.put(STORE_NAME, state)
110 } catch (error) {
111 console.error('Failed to save AI state to IndexedDB:', error)
112 }
113}
114
115async function clearStorage(): Promise<void> {
116 try {
117 const db = await openAiDb()
118 await db.clear(STORE_NAME)
119 } catch (error) {
120 console.error('Failed to clear AI state from IndexedDB:', error)
121 }
122}
123
124// Helper function to sanitize objects to ensure they're cloneable
125// Issue due to addToolResult
126function sanitizeForCloning(obj: any): any {
127 if (obj === null || obj === undefined) return obj
128 if (typeof obj !== 'object') return obj
129 return JSON.parse(JSON.stringify(obj))
130}
131
132// Helper function to load state from IndexedDB
133async function loadFromIndexedDB(projectRef: string): Promise<StoredAiAssistantState | null> {
134 try {
135 const persistedState = await getAiState(projectRef)
136 if (persistedState) {
137 // Revive dates and sanitize message data
138 Object.values(persistedState.chats).forEach((chat: ChatSession) => {
139 if (chat && typeof chat === 'object') {
140 chat.createdAt = new Date(chat.createdAt)
141 chat.updatedAt = new Date(chat.updatedAt)
142
143 // Sanitize message parts to remove proxy objects
144 if (chat.messages) {
145 chat.messages.forEach((message: any) => {
146 if (message.parts) {
147 message.parts = message.parts.map((part: any) => sanitizeForCloning(part))
148 }
149 })
150 }
151 }
152 })
153 return persistedState
154 }
155 } catch (error) {
156 console.error('Error loading AI state from IndexedDB:', error)
157 }
158 return null
159}
160
161// Helper function to attempt migration from localStorage
162async function tryMigrateFromLocalStorage(
163 projectRef: string
164): Promise<StoredAiAssistantState | null> {
165 const stored = localStorage.getItem(LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE(projectRef))
166 if (!stored) {
167 return null
168 }
169
170 let migratedState: StoredAiAssistantState | null = null
171 try {
172 const parsedFromLocalStorage = JSON.parse(stored, (key, value) => {
173 if ((key === 'createdAt' || key === 'updatedAt') && value) {
174 return new Date(value)
175 }
176 return value
177 })
178
179 if (parsedFromLocalStorage && typeof parsedFromLocalStorage.chats === 'object') {
180 migratedState = {
181 projectRef: projectRef,
182 activeChatId: parsedFromLocalStorage.activeChatId,
183 chats: parsedFromLocalStorage.chats,
184 model: parsedFromLocalStorage.model ?? INITIAL_AI_ASSISTANT.model,
185 }
186 } else {
187 console.warn('Data in localStorage is not in the expected format, ignoring.')
188 // Clean up invalid data
189 localStorage.removeItem(LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE(projectRef))
190 }
191 } catch (error) {
192 console.error('Failed to parse state from localStorage:', error)
193 // Clear potentially corrupted data
194 localStorage.removeItem(LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE(projectRef))
195 }
196
197 if (migratedState) {
198 try {
199 await saveAiState(migratedState)
200 localStorage.removeItem(LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE(projectRef))
201 return migratedState
202 } catch (saveError) {
203 console.error('Failed to save migrated state to IndexedDB:', saveError)
204 return null
205 }
206 }
207
208 return null
209}
210
211// Helper function to ensure an active chat exists or initialize a new one
212function ensureActiveChatOrInitialize(state: AiAssistantState) {
213 // Ensure an active chat exists after loading/migration
214 if (!state.activeChatId || !state.chats[state.activeChatId]) {
215 const chatIds = Object.keys(state.chats)
216 if (chatIds.length > 0) {
217 // Select the most recently updated chat
218 state.activeChatId = chatIds.sort(
219 (a, b) =>
220 (state.chats[b].updatedAt?.getTime() || 0) - (state.chats[a].updatedAt?.getTime() || 0)
221 )[0]
222 } else {
223 // If loaded/migrated state had no chats, create a new one
224 state.newChat()
225 }
226 }
227}
228
229function createChatInstance(
230 state: AiAssistantState,
231 options: { id: string; initialMessages: MessageType[] }
232) {
233 return new Chat<MessageType>({
234 id: options.id,
235 messages: options.initialMessages.map((message) => sanitizeForCloning(message)),
236 sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,
237 transport: new DefaultChatTransport({
238 api: `${BASE_PATH}/api/ai/sql/generate-v4`,
239 fetch: async (url, init) => {
240 const response = await globalThis.fetch(url as RequestInfo, init)
241 const spanId = response.headers.get('x-braintrust-span-id')
242 if (spanId) {
243 state.pendingSpanIds[options.id] = spanId
244 }
245 return response
246 },
247 async prepareSendMessagesRequest({ messages, ...opts }) {
248 const cleanedMessages = prepareMessagesForAPI(messages)
249 const headerData = await constructHeaders()
250 const authorizationHeader = headerData.get('Authorization')
251
252 // Get the chat specific to this request to ensure we have the correct name
253 const chat = state.chats[options.id]
254
255 return {
256 ...opts,
257 body: {
258 messages: cleanedMessages,
259 projectRef: state.context.projectRef,
260 connectionString: state.context.connectionString,
261 chatId: options.id,
262 chatName: chat?.name,
263 orgSlug: state.context.orgSlug,
264 context: state.context,
265 model: state.model,
266 ...opts.body,
267 },
268 ...(IS_PLATFORM ? { headers: { Authorization: authorizationHeader ?? '' } } : {}),
269 }
270 },
271 }),
272 async onToolCall({ toolCall }) {
273 if (toolCall.dynamic) {
274 return
275 }
276
277 if (toolCall.toolName === 'rename_chat') {
278 const { newName } = toolCall.input as { newName: string }
279
280 if (options.id && newName?.trim()) {
281 state.renameChat(options.id, newName.trim())
282 }
283 }
284 },
285 onFinish(_result) {
286 // Sync messages back to state
287 const chatInstance = state.chatInstances[options.id]
288 if (chatInstance) {
289 const messages = chatInstance.messages
290 const chat = state.chats[options.id]
291 if (chat) {
292 chat.messages = messages
293 chat.updatedAt = new Date()
294 }
295
296 // Associate pending span ID with the last assistant message
297 const pendingSpanId = state.pendingSpanIds[options.id]
298 if (pendingSpanId) {
299 const lastAssistantMsg = [...messages].reverse().find((m) => m.role === 'assistant')
300 if (lastAssistantMsg) {
301 state.messageSpanIds[lastAssistantMsg.id] = pendingSpanId
302 }
303 delete state.pendingSpanIds[options.id]
304 }
305 }
306 },
307 })
308}
309
310export const createAiAssistantState = (): AiAssistantState => {
311 // Initialize with defaults, loading happens asynchronously in the provider
312 const initialState = { ...INITIAL_AI_ASSISTANT }
313
314 const state: AiAssistantState = proxy({
315 ...initialState, // Spread initial values directly
316 chatInstances: {},
317 pendingSpanIds: {},
318 messageSpanIds: {},
319
320 setContext: (context: Partial<AiAssistantContext>) => {
321 state.context = { ...state.context, ...context }
322 },
323
324 resetAiAssistantPanel: () => {
325 Object.assign(state, INITIAL_AI_ASSISTANT)
326 },
327
328 setModel: (model: AssistantModel) => {
329 state.model = model
330 },
331
332 // Chat management
333 get activeChat(): ChatSession | undefined {
334 return state.activeChatId ? state.chats[state.activeChatId] : undefined
335 },
336
337 newChat: (
338 options?: { name?: string; initialMessage?: string } & Partial<
339 Pick<AiAssistantData, 'initialInput' | 'sqlSnippets' | 'suggestions' | 'tables'>
340 >
341 ) => {
342 const chatId = uuidv4()
343 const newChat: ChatSession = {
344 id: chatId,
345 name: options?.name ?? 'New chat',
346 messages: [],
347 createdAt: new Date(),
348 updatedAt: new Date(),
349 }
350
351 state.chats = {
352 ...state.chats,
353 [chatId]: newChat,
354 }
355 state.activeChatId = chatId
356
357 // Create new chat instance
358 const chatInstance = createChatInstance(state, { id: chatId, initialMessages: [] })
359
360 state.chatInstances[chatId] = ref(chatInstance)
361
362 // If initialMessage is provided, append it to the chat instance
363 if (options?.initialMessage) {
364 chatInstance.sendMessage({
365 text: options.initialMessage,
366 })
367 }
368
369 // Update non-chat related state based on options, falling back to current state, then initial
370 state.initialInput = options?.initialInput ?? INITIAL_AI_ASSISTANT.initialInput
371 state.sqlSnippets = options?.sqlSnippets ?? INITIAL_AI_ASSISTANT.sqlSnippets
372 state.suggestions = options?.suggestions ?? INITIAL_AI_ASSISTANT.suggestions
373 state.tables = options?.tables ?? INITIAL_AI_ASSISTANT.tables
374
375 return chatId
376 },
377
378 selectChat: (id: string) => {
379 if (id !== state.activeChatId) {
380 state.activeChatId = id
381 const chat = state.chats[id]
382 if (chat) {
383 if (!state.chatInstances[id]) {
384 state.chatInstances[id] = ref(
385 createChatInstance(state, { id, initialMessages: chat.messages })
386 )
387 }
388 }
389 }
390 },
391
392 deleteChat: (id: string) => {
393 const { [id]: _, ...remainingChats } = state.chats
394 state.chats = remainingChats
395
396 if (id === state.activeChatId) {
397 const remainingChatIds = Object.keys(remainingChats)
398 state.activeChatId = remainingChatIds.length > 0 ? remainingChatIds[0] : undefined
399
400 if (state.activeChatId) {
401 const chat = state.chats[state.activeChatId]
402 if (!state.chatInstances[state.activeChatId]) {
403 state.chatInstances[state.activeChatId] = ref(
404 createChatInstance(state, { id: state.activeChatId, initialMessages: chat.messages })
405 )
406 }
407 }
408 }
409 },
410
411 renameChat: (id: string, name: string) => {
412 const chat = state.chats[id]
413 if (chat && chat.name !== name) {
414 chat.name = name
415 chat.updatedAt = new Date()
416 }
417 },
418
419 clearMessages: () => {
420 const chat = state.activeChat
421 if (chat) {
422 chat.messages = []
423 chat.updatedAt = new Date()
424 state.suggestions = undefined
425 state.sqlSnippets = []
426 state.initialInput = ''
427 }
428 },
429
430 deleteMessagesAfter: (id: string, { includeSelf = true } = {}) => {
431 const chat = state.activeChat
432 if (!chat) return
433
434 const messageIndex = chat.messages.findIndex((msg) => msg.id === id)
435 if (messageIndex === -1) return
436
437 // Delete all messages from the target message (optionally including) to the end
438 const startIndex = includeSelf ? messageIndex : messageIndex + 1
439 chat.messages.splice(startIndex)
440 chat.updatedAt = new Date()
441 },
442
443 saveMessage: (message: MessageType | MessageType[]) => {
444 const chat = state.activeChat
445 if (!chat) return
446
447 const incomingMessages = Array.isArray(message) ? message : [message]
448
449 const messagesToAdd: AssistantMessageType[] = []
450
451 incomingMessages.forEach((msg) => {
452 const index = chat.messages.findIndex((existing) => existing.id === msg.id)
453
454 if (index !== -1) {
455 state.updateMessage(msg)
456 } else {
457 messagesToAdd.push(msg)
458 }
459 })
460
461 if (messagesToAdd.length > 0) {
462 chat.messages.push(...messagesToAdd)
463 chat.updatedAt = new Date()
464 }
465 },
466
467 updateMessage: (updatedMessage: MessageType) => {
468 const chat = state.activeChat
469 if (!chat) return
470
471 const messageIndex = chat.messages.findIndex((msg) => msg.id === updatedMessage.id)
472 if (messageIndex !== -1) {
473 chat.messages[messageIndex] = updatedMessage
474 chat.updatedAt = new Date()
475 }
476 },
477
478 setSqlSnippets: (snippets: SqlSnippet[]) => {
479 state.sqlSnippets = snippets
480 },
481
482 clearSqlSnippets: () => {
483 state.sqlSnippets = undefined
484 state.suggestions = undefined
485 },
486
487 // --- New function to load persisted state ---
488 loadPersistedState: (persistedState: StoredAiAssistantState) => {
489 state.chats = persistedState.chats
490 state.activeChatId = persistedState.activeChatId
491 const storedModel = persistedState.model
492 state.model =
493 storedModel && isKnownAssistantModelId(storedModel)
494 ? storedModel
495 : INITIAL_AI_ASSISTANT.model
496
497 // Ensure an active chat exists after loading
498 if (!state.activeChat) {
499 const chatIds = Object.keys(state.chats)
500 if (chatIds.length > 0) {
501 // Select the most recently updated chat
502 state.activeChatId = chatIds.sort(
503 (a, b) =>
504 (state.chats[b].updatedAt?.getTime() || 0) -
505 (state.chats[a].updatedAt?.getTime() || 0)
506 )[0]
507 } else {
508 // If loaded state had no chats, create a new one
509 state.newChat()
510 }
511 }
512
513 // Initialize chat instance for the active chat
514 if (
515 state.activeChatId &&
516 state.chats[state.activeChatId] &&
517 !state.chatInstances[state.activeChatId]
518 ) {
519 state.chatInstances[state.activeChatId] = ref(
520 createChatInstance(state, {
521 id: state.activeChatId,
522 initialMessages: state.chats[state.activeChatId].messages,
523 })
524 )
525 }
526 },
527
528 clearStorage: async () => {
529 await clearStorage()
530 },
531 })
532
533 return state
534}
535
536export type AiAssistantState = AiAssistantData & {
537 resetAiAssistantPanel: () => void
538 activeChat: ChatSession | undefined
539 chatInstances: Record<string, Chat<MessageType>>
540 pendingSpanIds: Record<string, string>
541 messageSpanIds: Record<string, string>
542 setContext: (context: Partial<AiAssistantContext>) => void
543 setModel: (model: AssistantModel) => void
544 newChat: (
545 options?: { name?: string; initialMessage?: string } & Partial<
546 Pick<AiAssistantData, 'initialInput' | 'sqlSnippets' | 'suggestions' | 'tables'>
547 >
548 ) => string
549 selectChat: (id: string) => void
550 deleteChat: (id: string) => void
551 renameChat: (id: string, name: string) => void
552 clearMessages: () => void
553 deleteMessagesAfter: (id: string, options?: { includeSelf?: boolean }) => void
554 saveMessage: (message: MessageType | MessageType[]) => void
555 updateMessage: (message: MessageType) => void
556 setSqlSnippets: (snippets: SqlSnippet[]) => void
557 clearSqlSnippets: () => void
558 loadPersistedState: (persistedState: StoredAiAssistantState) => void
559 clearStorage: () => Promise<void>
560}
561
562export const AiAssistantStateContext = createContext<AiAssistantState>(createAiAssistantState())
563
564export const AiAssistantStateContextProvider = ({ children }: PropsWithChildren) => {
565 const { data: project } = useSelectedProjectQuery()
566 // Initialize state. createAiAssistantState now just sets defaults.
567 const [state] = useState(() => createAiAssistantState())
568
569 // Effect to load state from IndexedDB on mount or projectRef change
570 useEffect(() => {
571 let isMounted = true
572
573 async function loadAndInitializeState() {
574 if (!project?.ref || typeof window === 'undefined') {
575 if (project?.ref === undefined) {
576 state.resetAiAssistantPanel()
577 }
578 return // Don't load if no projectRef or not in browser
579 }
580
581 let loadedState: StoredAiAssistantState | null = null
582
583 // 1. Try loading from IndexedDB
584 loadedState = await loadFromIndexedDB(project?.ref)
585
586 // 2. If not in IndexedDB, try migrating from localStorage
587 if (!loadedState) {
588 loadedState = await tryMigrateFromLocalStorage(project?.ref)
589 }
590
591 if (!isMounted) return // Component unmounted during async operations
592
593 // 3. If state was loaded or migrated, update the valtio state
594 if (loadedState) {
595 state.loadPersistedState(loadedState)
596 }
597
598 // 4. Ensure an active chat exists and handle URL overrides
599 ensureActiveChatOrInitialize(state)
600 }
601
602 loadAndInitializeState()
603
604 return () => {
605 isMounted = false
606 }
607 }, [project?.ref, state])
608
609 // Effect to save state to IndexedDB on changes
610 useEffect(() => {
611 if (typeof window !== 'undefined' && project?.ref) {
612 // Create a debounced version of saveAiState
613 const debouncedSaveAiState = debounce(saveAiState, 500)
614
615 const unsubscribe = subscribe(state, () => {
616 const snap = snapshot(state)
617 // Prepare state for IndexedDB
618 const stateToSave: StoredAiAssistantState = {
619 projectRef: project?.ref,
620 activeChatId: snap.activeChatId,
621 model: snap.model,
622 chats: snap.chats
623 ? Object.entries(snap.chats).reduce((acc, [chatId, chat]) => {
624 // Limit messages before saving
625 return {
626 ...acc,
627 [chatId]: {
628 ...chat,
629 messages: chat.messages?.slice(-20) || [],
630 },
631 }
632 }, {})
633 : {},
634 }
635 debouncedSaveAiState(stateToSave)
636 })
637 // Clean up subscription and cancel any pending saves on unmount or projectRef change
638 return () => {
639 debouncedSaveAiState.cancel()
640 unsubscribe()
641 }
642 }
643 return undefined
644 }, [state, project?.ref])
645
646 return (
647 <AiAssistantStateContext.Provider value={state}>{children}</AiAssistantStateContext.Provider>
648 )
649}
650
651export const useAiAssistantStateSnapshot = (options?: Parameters<typeof useSnapshot>[1]) => {
652 const state = useContext(AiAssistantStateContext)
653 return useSnapshot(state, options)
654}
655
656export const useAiAssistantState = () => {
657 const state = useContext(AiAssistantStateContext)
658 return state
659}