Spaces:
Running
Running
File size: 2,289 Bytes
84c9f51 a86b547 84c9f51 a86b547 84c9f51 a86b547 84c9f51 a86b547 84c9f51 a86b547 84c9f51 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import { useChat, type Message, UseChatHelpers } from 'ai/react';
import { toast } from 'react-hot-toast';
import { useEffect, useState } from 'react';
import { ChatEntity, MessageBase } from '../types';
import { saveKVChatMessage } from '../kv/chat';
const useVisionAgent = (chat: ChatEntity) => {
const { messages: initialMessages, id, url } = chat;
const {
messages,
append: appendRaw,
reload,
stop,
isLoading,
input,
setInput,
setMessages,
} = useChat({
sendExtraMessageFields: true,
api: '/api/vision-agent',
onResponse(response) {
if (response.status !== 200) {
toast.error(response.statusText);
}
},
onFinish(message) {
saveKVChatMessage(id, message);
},
initialMessages: initialMessages,
body: {
url,
id,
},
});
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]);
useEffect(() => {
if (
!isLoading &&
messages.length &&
messages[messages.length - 1].role === 'user'
) {
reload();
}
}, [isLoading, messages, reload]);
const assistantLoadingMessage = {
id: 'loading',
content: loadingDots,
role: 'assistant',
};
const messageWithLoading =
isLoading &&
messages.length &&
messages[messages.length - 1].role !== 'assistant'
? [...messages, assistantLoadingMessage]
: messages;
const append: UseChatHelpers['append'] = async message => {
await saveKVChatMessage(id, message as MessageBase);
return appendRaw(message);
};
return {
messages: messageWithLoading as MessageBase[],
append,
reload,
stop,
isLoading,
input,
setInput,
};
};
export default useVisionAgent;
|