import { useChat, type Message, UseChatHelpers } from 'ai/react'; import { toast } from 'react-hot-toast'; import { useEffect, useState } from 'react'; import { MessageBase, SignedPayload } from '../types'; import { fetcher, nanoid } from '../utils'; import { getCleanedUpMessages, generateAnswersImageMarkdown, generateInputImageMarkdown, } from '../messageUtils'; import { CLEANED_SEPARATOR } from '../constants'; import { useSearchParams } from 'next/navigation'; import { ChatWithMessages, MessageRaw } from '../db/types'; import { dbPostCreateMessage } from '../db/functions'; const uploadBase64 = async ( base64: string, messageId: string, chatId: string, index: number, ) => { const res = await fetch( 'data:image/png;base64,' + base64.replace('base:64', ''), ); const blob = await res.blob(); const { signedUrl, publicUrl, fields } = await fetcher( '/api/sign', { method: 'POST', body: JSON.stringify({ id: `${chatId}/${messageId}`, fileType: blob.type, fileName: `answer-${index}.${blob.type.split('/')[1]}`, }), }, ); const formData = new FormData(); Object.entries(fields).forEach(([key, value]) => { formData.append(key, value as string); }); formData.append('file', blob); const uploadResponse = await fetch(signedUrl, { method: 'POST', body: formData, }); if (uploadResponse.ok) { return publicUrl; } else { throw new Error('Upload failed'); } }; const useVisionAgent = (chat: ChatWithMessages) => { const { messages: initialMessages, id, mediaUrl } = chat; const { messages, append: appendRaw, isLoading, reload, } = useChat({ api: '/api/vision-agent', // @ts-ignore https://sdk.vercel.ai/docs/troubleshooting/common-issues/use-chat-failed-to-parse-stream streamMode: 'text', onResponse(response) { if (response.status !== 200) { toast.error(response.statusText); } }, onFinish: async message => { await dbPostCreateMessage(id, { role: message.role as 'user' | 'assistant', content: message.content, }); }, initialMessages: initialMessages, body: { mediaUrl, id, }, }); /** * If case this is first time user navigated with init message, we need to reload the chat for the first response */ useEffect(() => { if (!isLoading && messages.length === 1 && messages[0].role === 'user') { reload(); } }, [isLoading, messages, reload]); const append: UseChatHelpers['append'] = async message => { dbPostCreateMessage(id, { role: message.role as 'user' | 'assistant', content: message.content, }); return appendRaw(message); }; return { messages: messages as MessageBase[], append, reload, isLoading, }; }; export default useVisionAgent;