'use client'; import { useState, useCallback, useRef, useEffect } from 'react'; import { useChat } from '@ai-sdk/react'; import { DefaultChatTransport, type ChatStatus, type UIMessage } from 'ai'; import { DEFAULT_OPENROUTER_MODEL_ID, type OpenRouterModelId } from '@/lib/openrouter'; import type { PromptInputMessage } from '@/components/ai-elements/prompt-input'; import { createChat, deleteChat, getMessages, listChats, saveMessages, type StoredChat, } from '@/services/indexed-db/chat-store'; async function convertToDataUrl(url: string): Promise { const response = await fetch(url); if (!response.ok) throw new Error('Could not read the selected file'); const blob = await response.blob(); return new Promise((resolve, reject) => { const reader = new FileReader(); reader.onload = () => { if (typeof reader.result === 'string') { resolve(reader.result); } else { reject(new Error('Could not encode the selected file')); } }; reader.onerror = () => reject(reader.error); reader.readAsDataURL(blob); }); } export function useChatSession(defaultModelId: OpenRouterModelId = DEFAULT_OPENROUTER_MODEL_ID) { const [modelId, setModelId] = useState(defaultModelId); const [inputValue, setInputValue] = useState(''); const [chatId, setChatId] = useState(null); const [chats, setChats] = useState([]); const chatIdRef = useRef(chatId); const pendingChatIdRef = useRef(null); chatIdRef.current = chatId; const fetchChats = useCallback(async () => { try { setChats(await listChats()); } catch (error) { console.error('Failed to fetch chats:', error); } }, []); const transportRef = useRef( new DefaultChatTransport({ api: '/api/chat', }) ); const { status, messages: chatMessages, setMessages: setChatMessages, sendMessage, stop, error: aiError, } = useChat({ transport: transportRef.current, onFinish: async ({ messages }) => { const activeChatId = pendingChatIdRef.current ?? chatIdRef.current; if (!activeChatId) return; try { await saveMessages(activeChatId, messages); await fetchChats(); } catch (error) { console.error('Failed to save messages:', error); } finally { if (pendingChatIdRef.current === activeChatId) { pendingChatIdRef.current = null; } } }, onError: (error) => { console.error('Chat error:', error); }, }); const [isPreparingFiles, setIsPreparingFiles] = useState(false); const chatStatus: ChatStatus = isPreparingFiles ? 'submitted' : status === 'submitted' || status === 'streaming' ? status : 'ready'; useEffect(() => { fetchChats(); }, [fetchChats]); const handleSelectChat = useCallback( async (selectedChatId: string) => { try { const messages = await getMessages(selectedChatId); chatIdRef.current = selectedChatId; setChatId(selectedChatId); setChatMessages(messages); } catch (error) { console.error('Failed to load chat:', error); } }, [setChatMessages] ); const handleNewChat = useCallback(() => { chatIdRef.current = null; setChatId(null); setChatMessages([]); setInputValue(''); }, [setChatMessages]); const handleDeleteChat = useCallback( async (deletedChatId: string) => { try { await deleteChat(deletedChatId); setChats((prev) => prev.filter((c) => c.id !== deletedChatId)); if (chatId === deletedChatId) { chatIdRef.current = null; setChatId(null); setChatMessages([]); } } catch (error) { console.error('Failed to delete chat:', error); } }, [chatId, setChatMessages] ); const handleSendMessage = useCallback( async (message: PromptInputMessage) => { const trimmedValue = message.text?.trim(); if (!trimmedValue && !message.files?.length) return; if ( status === 'submitted' || status === 'streaming' || isPreparingFiles ) return; let activeChatId = chatIdRef.current; if (!chatIdRef.current) { const newChatId = crypto.randomUUID(); const title = trimmedValue ? trimmedValue.slice(0, 100) : 'New Chat'; activeChatId = newChatId; chatIdRef.current = newChatId; setChatId(newChatId); try { await createChat(newChatId, title); await fetchChats(); } catch (error) { console.error('Failed to create chat:', error); } } const parts: Array< | { type: 'text'; text: string } | { type: 'file'; url: string; mediaType: string; filename: string } > = []; if (trimmedValue) { parts.push({ type: 'text', text: trimmedValue }); } // PromptInput supplies files as base64 data URLs. Keep them inline so the // demo never needs a remote storage bucket. if (message.files?.length) { setIsPreparingFiles(true); for (const file of message.files) { try { const dataUrl = file.url.startsWith('data:') ? file.url : await convertToDataUrl(file.url); parts.push({ type: 'file', url: dataUrl, mediaType: file.mediaType || 'application/octet-stream', filename: file.filename || 'file', }); } catch (err) { console.error('File encoding error:', err); } } setIsPreparingFiles(false); } if (parts.length === 0) return; const userMessage: UIMessage = { id: crypto.randomUUID(), role: 'user', parts, }; setInputValue(''); if (activeChatId) { try { await saveMessages(activeChatId, [...chatMessages, userMessage]); } catch (error) { console.error('Failed to save user message:', error); } } pendingChatIdRef.current = activeChatId; await sendMessage(userMessage, { body: { modelId } }); }, [ status, sendMessage, chatMessages, isPreparingFiles, fetchChats, modelId, ] ); const handleStop = useCallback(() => { stop(); }, [stop]); return { inputValue, setInputValue, modelId, setModelId, chatId, chats, chatMessages, status, chatStatus, aiError, isPreparingFiles, handleSelectChat, handleNewChat, handleDeleteChat, handleSendMessage, handleStop, }; }