Spaces:
Sleeping
Sleeping
File size: 3,254 Bytes
84c9f51 a86b547 5ec491a 96ac62a 0d6f04b 0a08909 d553ae5 5ec491a 5d7d435 0a08909 5d7d435 a86b547 5ec491a d553ae5 a86b547 84c9f51 a86b547 314f2dc a86b547 5d7d435 3ebf44a 84c9f51 a86b547 5ec491a a86b547 d553ae5 a86b547 5ec491a 84c9f51 5ec491a 84c9f51 a86b547 5ec491a 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
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<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: ChatWithMessages) => {
const { messages: initialMessages, id, mediaUrl } = chat;
const searchParams = useSearchParams();
const reflectionValue = searchParams.get('reflection');
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 => {
await dbPostCreateMessage(id, {
role: message.role as 'user' | 'assistant',
content: message.content,
});
},
initialMessages: initialMessages,
body: {
mediaUrl,
id,
enableSelfReflection: reflectionValue === 'true',
},
});
/**
* If the last message is from the user, reload the chat, this would trigger to get the response from the assistant
* There are 2 scenarios when this might happen
* 1. Navigated from example images, init message only include preset user message
* 2. Last time the assistant message failed or not saved to database.
*/
useEffect(() => {
if (
!isLoading &&
messages.length &&
messages[messages.length - 1].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,
stop,
isLoading,
input,
setInput,
};
};
export default useVisionAgent;
|