Spaces:
Sleeping
Sleeping
File size: 2,373 Bytes
a1c5622 a86b547 a1c5622 0fd8446 ca7a659 a1c5622 4af6326 ae074fc a86b547 5ec491a a1c5622 4af6326 d553ae5 a1c5622 a86b547 92f037b a86b547 5d7d435 a1c5622 4af6326 84c9f51 a1c5622 a86b547 a1c5622 a86b547 ca7a659 a86b547 5ec491a 92f037b 5ec491a c232e44 84c9f51 c232e44 84c9f51 a86b547 0fd8446 ae074fc a1c5622 ae074fc a1c5622 a86b547 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import { useChat } from 'ai/react';
import { toast } from 'react-hot-toast';
import { useEffect, useRef, useState } from 'react';
import { ChatWithMessages, MessageUI, MessageUserInput } from '../types';
import {
dbPostCreateMessage,
dbPostUpdateMessageResponse,
} from '../db/functions';
import {
convertAssistantUIMessageToDBMessageResponse,
convertDBMessageToUIMessage,
} from '../utils/message';
import { useSetAtom } from 'jotai';
import { selectedMessageId } from '@/state/chat';
import { Message } from '@prisma/client';
const useVisionAgent = (chat: ChatWithMessages) => {
const { messages: dbMessages, id, mediaUrl } = chat;
const latestDbMessage = dbMessages[dbMessages.length - 1];
const setMessageId = useSetAtom(selectedMessageId);
// Temporary solution for now while single we have to pass mediaUrl separately outside of the messages
const currMediaUrl = useRef<string>(mediaUrl);
const currMessageId = useRef<string>(latestDbMessage?.id);
const { messages, append, isLoading, reload } = useChat({
api: '/api/vision-agent',
streamMode: 'text',
onResponse(response) {
if (response.status !== 200) {
toast.error(response.statusText);
}
},
onFinish: async message => {
await dbPostUpdateMessageResponse(
currMessageId.current,
convertAssistantUIMessageToDBMessageResponse(message),
);
setMessageId(currMessageId.current);
},
initialMessages: convertDBMessageToUIMessage(dbMessages),
body: {
mediaUrl: currMediaUrl.current,
id,
},
onError: err => {
err && toast.error(err.message);
},
});
/**
* If case this is first time user navigated with init message, we need to reload the chat for the first response
*/
const once = useRef(true);
useEffect(() => {
if (
!isLoading &&
messages.length === 1 &&
messages[0].role === 'user' &&
once.current
) {
once.current = false;
reload();
}
}, [isLoading, messages, reload]);
return {
messages: messages as MessageUI[],
append: (message: Message) => {
currMediaUrl.current = message.mediaUrl;
currMessageId.current = message.id;
append({
id,
role: 'user',
content: message.prompt,
});
},
reload,
isLoading,
};
};
export default useVisionAgent;
|