import { useChat, type Message } from 'ai/react'; import { toast } from 'react-hot-toast'; import { useEffect, useState } from 'react'; import { MediaDetails } from '../fetch'; import { MessageWithSelectedDataset } from '../types'; const useChatWithMedia = (mediaList: MediaDetails[]) => { const { messages, append, reload, stop, isLoading, input, setInput, setMessages, } = useChat({ sendExtraMessageFields: true, onResponse(response) { if (response.status !== 200) { toast.error(response.statusText); } }, initialMessages: [ { id: 'system', content: `For the full conversation, user have provided the following images: ${mediaList.map(media => media.name)}. Please help reply to user regarding these images`, dataset: mediaList, role: 'system', }, ] as MessageWithSelectedDataset[], }); const [loadingDots, setLoadingDots] = useState(''); useEffect(() => { let loadingInterval: NodeJS.Timeout; if (isLoading) { loadingInterval = setInterval(() => { setLoadingDots(prevMessage => { switch (prevMessage) { case '': return '.'; case '.': return '..'; case '..': return '...'; case '...': return ''; default: return ''; } }); }, 500); } return () => { clearInterval(loadingInterval); }; }, [isLoading]); const assistantLoadingMessage = { id: 'loading', content: loadingDots, role: 'assistant', }; const messageWithLoading = isLoading && messages.length && messages[messages.length - 1].role !== 'assistant' ? [...messages, assistantLoadingMessage] : messages; return { messages: messageWithLoading as MessageWithSelectedDataset[], append, reload, stop, isLoading, input, setInput, }; }; export default useChatWithMedia;