Spaces:
Sleeping
Sleeping
import { useChat, type Message, UseChatHelpers } from 'ai/react'; | |
import { toast } from 'react-hot-toast'; | |
import { useEffect, useState } from 'react'; | |
import { ChatEntity, MessageBase, SignedPayload } from '../types'; | |
import { saveKVChatMessage } from '../kv/chat'; | |
import { fetcher, nanoid } from '../utils'; | |
import { | |
getCleanedUpMessages, | |
generateAnswersImageMarkdown, | |
generateInputImageMarkdown, | |
} from '../messageUtils'; | |
import { CLEANED_SEPARATOR } from '../constants'; | |
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<SignedPayload>( | |
'/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: ChatEntity, enableSelfReflection = true) => { | |
const { messages: initialMessages, id, url } = chat; | |
const { | |
messages, | |
append: appendRaw, | |
reload, | |
stop, | |
isLoading, | |
input, | |
setInput, | |
setMessages, | |
error, | |
} = useChat({ | |
api: '/api/vision-agent', | |
onResponse(response) { | |
if (response.status !== 200) { | |
toast.error(response.statusText); | |
} | |
}, | |
onFinish: async message => { | |
const { logs = '', content, images } = getCleanedUpMessages(message); | |
if (images?.length) { | |
const publicUrls = await Promise.all( | |
images.map((image, index) => | |
uploadBase64(image, message.id, id ?? 'no-id', index), | |
), | |
); | |
const newContent = publicUrls.reduce((accum, url, index) => { | |
return accum.replace( | |
generateAnswersImageMarkdown(index, '/loading.gif'), | |
generateAnswersImageMarkdown(index, url), | |
); | |
}, content); | |
const newMessage = { | |
...message, | |
content: logs + CLEANED_SEPARATOR + newContent, | |
}; | |
setMessages([ | |
...messages, | |
/** | |
* A workaround to fix the issue of the messages been stale state when appending a new message | |
* https://github.com/vercel/ai/issues/550#issuecomment-1712693371 | |
*/ | |
...(input | |
? [ | |
{ | |
id: nanoid(), | |
role: 'user', | |
content: input + '\n\n' + generateInputImageMarkdown(url), | |
createdAt: new Date(), | |
} satisfies Message, | |
] | |
: []), | |
newMessage, | |
]); | |
if (id) { | |
saveKVChatMessage(id, newMessage); | |
} | |
} else { | |
if (id) { | |
saveKVChatMessage(id, { | |
...message, | |
content: logs + CLEANED_SEPARATOR + content, | |
}); | |
} | |
} | |
}, | |
initialMessages: initialMessages, | |
body: { | |
url, | |
id, | |
enableSelfReflection, | |
}, | |
}); | |
useEffect(() => { | |
if ( | |
!isLoading && | |
messages.length && | |
messages[messages.length - 1].role === 'user' | |
) { | |
reload(); | |
} | |
}, [isLoading, messages, reload]); | |
const assistantLoadingMessage = { | |
id: 'loading', | |
content: '...', | |
role: 'assistant', | |
}; | |
const messageWithLoading = | |
isLoading && | |
messages.length && | |
messages[messages.length - 1].role !== 'assistant' | |
? [...messages, assistantLoadingMessage] | |
: messages; | |
const append: UseChatHelpers['append'] = async message => { | |
if (id) { | |
await saveKVChatMessage(id, message as MessageBase); | |
} | |
return appendRaw(message); | |
}; | |
return { | |
messages: messageWithLoading as MessageBase[], | |
append, | |
reload, | |
stop, | |
isLoading, | |
input, | |
setInput, | |
}; | |
}; | |
export default useVisionAgent; | |