@@ -57,7 +57,6 @@ export const AIAssistantHeader = ({
icon={
}
onClick={onNewChat}
className="h-7 w-7 p-0"
- disabled={isChatLoading}
tooltip={{ content: { side: 'bottom', text: 'New chat' } }}
/>
activeChatId?: string
model: AssistantModel
+ context: AiAssistantContext
}
// Data structure stored in IndexedDB
@@ -53,6 +65,7 @@ const INITIAL_AI_ASSISTANT: AiAssistantData = {
chats: {},
activeChatId: undefined,
model: 'gpt-5',
+ context: {},
}
const DB_NAME = 'ai-assistant-db'
@@ -211,12 +224,79 @@ function ensureActiveChatOrInitialize(state: AiAssistantState) {
}
}
+function createChatInstance(
+ state: AiAssistantState,
+ options: { id: string; initialMessages: MessageType[] }
+) {
+ return new Chat({
+ id: options.id,
+ messages: options.initialMessages,
+ sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
+ transport: new DefaultChatTransport({
+ api: `${BASE_PATH}/api/ai/sql/generate-v4`,
+ async prepareSendMessagesRequest({ messages, ...opts }) {
+ const cleanedMessages = prepareMessagesForAPI(messages)
+ const headerData = await constructHeaders()
+ const authorizationHeader = headerData.get('Authorization')
+
+ // Get the chat specific to this request to ensure we have the correct name
+ const chat = state.chats[options.id]
+
+ return {
+ ...opts,
+ body: {
+ messages: cleanedMessages,
+ projectRef: state.context.projectRef,
+ connectionString: state.context.connectionString,
+ chatName: chat?.name,
+ orgSlug: state.context.orgSlug,
+ context: state.context,
+ model: state.model,
+ ...opts.body,
+ },
+ ...(IS_PLATFORM ? { headers: { Authorization: authorizationHeader ?? '' } } : {}),
+ }
+ },
+ }),
+ async onToolCall({ toolCall }) {
+ if (toolCall.dynamic) {
+ return
+ }
+
+ if (toolCall.toolName === 'rename_chat') {
+ const { newName } = toolCall.input as { newName: string }
+
+ if (options.id && newName?.trim()) {
+ state.renameChat(options.id, newName.trim())
+ }
+ }
+ },
+ onFinish(result) {
+ // Sync messages back to state
+ const chatInstance = state.chatInstances[options.id]
+ if (chatInstance) {
+ const messages = chatInstance.messages
+ const chat = state.chats[options.id]
+ if (chat) {
+ chat.messages = messages as AssistantMessageType[]
+ chat.updatedAt = new Date()
+ }
+ }
+ },
+ })
+}
+
export const createAiAssistantState = (): AiAssistantState => {
// Initialize with defaults, loading happens asynchronously in the provider
const initialState = { ...INITIAL_AI_ASSISTANT }
const state: AiAssistantState = proxy({
...initialState, // Spread initial values directly
+ chatInstances: {},
+
+ setContext: (context: Partial) => {
+ state.context = { ...state.context, ...context }
+ },
resetAiAssistantPanel: () => {
Object.assign(state, INITIAL_AI_ASSISTANT)
@@ -232,7 +312,7 @@ export const createAiAssistantState = (): AiAssistantState => {
},
newChat: (
- options?: { name?: string } & Partial<
+ options?: { name?: string; initialMessage?: string } & Partial<
Pick
>
) => {
@@ -251,6 +331,18 @@ export const createAiAssistantState = (): AiAssistantState => {
}
state.activeChatId = chatId
+ // Create new chat instance
+ const chatInstance = createChatInstance(state, { id: chatId, initialMessages: [] })
+
+ state.chatInstances[chatId] = ref(chatInstance)
+
+ // If initialMessage is provided, append it to the chat instance
+ if (options?.initialMessage) {
+ chatInstance.sendMessage({
+ text: options.initialMessage,
+ })
+ }
+
// Update non-chat related state based on options, falling back to current state, then initial
state.initialInput = options?.initialInput ?? INITIAL_AI_ASSISTANT.initialInput
state.sqlSnippets = options?.sqlSnippets ?? INITIAL_AI_ASSISTANT.sqlSnippets
@@ -263,6 +355,14 @@ export const createAiAssistantState = (): AiAssistantState => {
selectChat: (id: string) => {
if (id !== state.activeChatId) {
state.activeChatId = id
+ const chat = state.chats[id]
+ if (chat) {
+ if (!state.chatInstances[id]) {
+ state.chatInstances[id] = ref(
+ createChatInstance(state, { id, initialMessages: chat.messages })
+ )
+ }
+ }
}
},
@@ -273,6 +373,15 @@ export const createAiAssistantState = (): AiAssistantState => {
if (id === state.activeChatId) {
const remainingChatIds = Object.keys(remainingChats)
state.activeChatId = remainingChatIds.length > 0 ? remainingChatIds[0] : undefined
+
+ if (state.activeChatId) {
+ const chat = state.chats[state.activeChatId]
+ if (!state.chatInstances[state.activeChatId]) {
+ state.chatInstances[state.activeChatId] = ref(
+ createChatInstance(state, { id: state.activeChatId, initialMessages: chat.messages })
+ )
+ }
+ }
}
},
@@ -382,6 +491,20 @@ export const createAiAssistantState = (): AiAssistantState => {
state.newChat()
}
}
+
+ // Initialize chat instance for the active chat
+ if (
+ state.activeChatId &&
+ state.chats[state.activeChatId] &&
+ !state.chatInstances[state.activeChatId]
+ ) {
+ state.chatInstances[state.activeChatId] = ref(
+ createChatInstance(state, {
+ id: state.activeChatId,
+ initialMessages: state.chats[state.activeChatId].messages,
+ })
+ )
+ }
},
clearStorage: async () => {
@@ -395,9 +518,11 @@ export const createAiAssistantState = (): AiAssistantState => {
export type AiAssistantState = AiAssistantData & {
resetAiAssistantPanel: () => void
activeChat: ChatSession | undefined
+ chatInstances: Record>
+ setContext: (context: Partial) => void
setModel: (model: AssistantModel) => void
newChat: (
- options?: { name?: string } & Partial<
+ options?: { name?: string; initialMessage?: string } & Partial<
Pick
>
) => string
@@ -508,3 +633,8 @@ export const useAiAssistantStateSnapshot = (options?: Parameters {
+ const state = useContext(AiAssistantStateContext)
+ return state
+}