Skip to content

Commit

Permalink
Merge pull request #119 from polyfire-ai/fix/usechat
Browse files Browse the repository at this point in the history 
fix: useChat hook
  • Loading branch information
@lowczarc
lowczarc authored Dec 12, 2023
2 parents 3341a60 + df85a32 commit f4e3bc1
Showing 1 changed file with 73 additions and 58 deletions.
131 changes: 73 additions & 58 deletions lib/hooks/useChat.ts
View file
Original file line number Diff line number Diff line change
@@ -1,29 +1,29 @@
/* eslint-disable camelcase */

import { useState, useEffect, useCallback } from "react";

import { generateUUIDV4 } from "../helpers/uuid";
import type { Chat as ChatType, ChatOptions } from "../chats";

import type { Chat, ChatOptions } from "../chats";
import usePolyfire from "./usePolyfire";

export type Message = {
id: string | null;
chat_id: string;
is_user_message: boolean;
content: string;
created_at: string | null;
end_of_message?: boolean;
id: string | null;
is_user_message: boolean;
};

export type ChatInfos = {
created_at: string;
id: string;
latest_message_content: string;
latest_message_created_at: string;
name: string | null;
system_prompt: string | null;
system_prompt_id: string | null;
created_at: string;
user_id: string;
latest_message_content: string;
latest_message_created_at: string;
};

const getChatList = async (getToken: () => Promise<string>): Promise<ChatInfos[]> => {
Expand Down Expand Up @@ -77,28 +77,28 @@ const renameChat = async (
};

type ChatBase<T> = {
loading: boolean;
error: string | undefined;
data: T | undefined;
error: string | undefined;
loading: boolean;
};

export type Chats = ChatBase<ChatInfos[]>;
export type ChatHistory = ChatBase<Message[]>;
export type ChatAnswer = ChatBase<Message>;

export type ChatUtils = {
onSendMessage: (message: string) => Promise<void>;
onSelectChat: (chatId: string) => Promise<void>;
onDeleteChat: (chatId: string) => Promise<void>;
onRenameChat: (chatId: string, name: string) => Promise<void>;
onResetChat: () => void;
onSelectChat: (chatId: string) => Promise<void>;
onSendMessage: (message: string) => Promise<void>;
};

export type ChatInstance = {
chats: Chats;
answer: ChatAnswer;
chat: ChatInfos | undefined;
chats: Chats;
history: ChatHistory;
answer: ChatAnswer;
utils: ChatUtils;
};

Expand All @@ -115,7 +115,7 @@ export default function useChat(
utils: { Chat },
} = usePolyfire();

const [chatInstance, setChatInstance] = useState<Chat>();
const [chatInstance, setChatInstance] = useState<ChatType>();
const [chatId, setChatId] = useState<string>();
// Chats list
const [chatsData, setChatsData] = useState<ChatInfos[]>([]);
Expand Down Expand Up @@ -157,14 +157,20 @@ export default function useChat(

const retrieveChat = useCallback(async (id: string) => {
setHistoryLoading(true);
const chatInstance = new Chat({ ...options, chatId: id } as ChatOptions);
const prevchatInstance = new Chat({
...options,
chatId: id,
} as ChatOptions);

setChatInstance(chatInstance);
setChatInstance(prevchatInstance);
setChatId(id);

chatInstance
prevchatInstance
.getMessages()
.then(setHistory)
.then((res) => {
if (res?.length) setHistory(res.reverse());
else setHistory([]);
})
.catch(setHistoryError)
.finally(() => {
setHistoryLoading(false);
Expand All @@ -175,53 +181,55 @@ export default function useChat(
async (id: string) => {
if (chatId === id) return;

const chat = chatsData.find((chat) => chat.id === id);
const chat = chatsData.find((c) => c.id === id);

if (!chat) return;

setChatData(chatsData.find((chat) => chat.id === id));
setChatData(chatsData.find((c) => c.id === id));
retrieveChat(id);
},
[chatsData, chatId],
);

const onCreateChat = useCallback(async (message: string) => {
// eslint-disable-next-line no-async-promise-executor
return new Promise<Chat>(async (resolve) => {
const newChatInstance = new Chat(options as ChatOptions);
setChatInstance(newChatInstance);

const newChatId = await newChatInstance.chatId;
setChatId(newChatId);

const newChatInfos = {
id: newChatId,
name: message,
system_prompt: newChatInstance?.options?.systemPrompt || null,
system_prompt_id: newChatInstance?.options?.systemPromptId || null,
created_at: new Date().toISOString(),
user_id: "",
latest_message_content: message,
latest_message_created_at: new Date().toISOString(),
};
const onCreateChat = useCallback(
async (message: string) =>
// eslint-disable-next-line no-async-promise-executor
new Promise<ChatType>(async (resolve) => {
const newChatInstance = new Chat(options as ChatOptions);
setChatInstance(newChatInstance);

const newChatId = await newChatInstance.chatId;
setChatId(newChatId);

const newChatInfos = {
id: newChatId,
name: message,
system_prompt: newChatInstance?.options?.systemPrompt || null,
system_prompt_id: newChatInstance?.options?.systemPromptId || null,
created_at: new Date().toISOString(),
user_id: "",
latest_message_content: message,
latest_message_created_at: new Date().toISOString(),
};

setChatsData((prev) => [newChatInfos, ...prev]);
setChatsData((prev) => [newChatInfos, ...prev]);

await onSelectChat(newChatId);
await onSelectChat(newChatId);

setHistory([]);
setHistoryError(undefined);
setHistory([]);
setHistoryError(undefined);

setAnswer(undefined);
setAnswerError(undefined);
setAnswer(undefined);
setAnswerError(undefined);

resolve(newChatInstance);
});
}, []);
resolve(newChatInstance);
}),
[],
);

const onRenameChat = useCallback(
async (id: string, name: string): Promise<void> => {
const chat = chatsData.find((chat) => chat.id === id);
const chat = chatsData.find((c) => c.id === id);

if (!chat) return;

Expand All @@ -234,16 +242,15 @@ export default function useChat(
);

const onDeleteChat = useCallback(
async (id: string): Promise<void> => {
return deleteChat(id, getToken).then(() => {
async (id: string): Promise<void> =>
deleteChat(id, getToken).then(() => {
if (chatId === id) {
setChatId(undefined);
setHistory([]);
setChatInstance(undefined);
}
setChatsData(chatsData.filter((chat) => chat.id !== id));
});
},
}),
[chatsData, chatId],
);

Expand Down Expand Up @@ -272,40 +279,48 @@ export default function useChat(
is_user_message: true,
content: message,
created_at: new Date().getTime().toString(),
end_of_message: true,
};
const aiMessage: Message = {
id: generateUUIDV4(),
chat_id: await newChatInstance.chatId,
is_user_message: false,
content: "",
created_at: null,
end_of_message: false,
};

setHistory((prev) => [userMessage, ...prev]);
setHistory((prev) => [...prev, userMessage]);
setAnswerError(undefined);

setAnswerLoading(true);

const stream = newChatInstance.sendMessage(message);

stream.on("data", (chunk: string) => {
aiMessage.content += chunk;

setAnswer({ ...aiMessage });
setAnswerLoading(false);
});

stream.on("error", (error: string) => {
console.error({ error });
setAnswerError(error);
onError?.(error);

stream.stop();
});

stream.on("end", async () => {
setAnswer(undefined);
aiMessage.created_at = new Date().getTime().toString();
setHistory((prev) => [aiMessage, ...prev]);
aiMessage.end_of_message = true;

setHistory((prev) => [...prev, aiMessage]);
setAnswer(undefined);

onSuccess?.();
});
} catch (error) {
console.error(error);
if (error instanceof Error) {
setAnswerError(error.message);
onError?.(error.message);
Expand Down

0 comments on commit f4e3bc1

Please sign in to comment.