From b73cc3b0df1d648555f5616ceee29dcb2d9d7de6 Mon Sep 17 00:00:00 2001 From: Srujan Gurram Date: Tue, 19 Sep 2023 21:23:31 +0530 Subject: [PATCH] Refactor code to handle chat completion and update chat messages and storage in real-time --- src/components/Sidebar/chat/chatInput.tsx | 5 +- src/components/Sidebar/chat/index.tsx | 14 +++--- src/hooks/useChatCompletion.tsx | 57 ++++++++++++++++++----- src/hooks/useCurrentChat.tsx | 55 ++++++++++++++++++++++ src/hooks/useHistory.tsx | 4 +- 5 files changed, 113 insertions(+), 22 deletions(-) create mode 100644 src/hooks/useCurrentChat.tsx diff --git a/src/components/Sidebar/chat/chatInput.tsx b/src/components/Sidebar/chat/chatInput.tsx index 638ffa79..a378e17c 100644 --- a/src/components/Sidebar/chat/chatInput.tsx +++ b/src/components/Sidebar/chat/chatInput.tsx @@ -3,11 +3,10 @@ import TextareaAutosize from 'react-textarea-autosize' import { GiMagicBroom } from 'react-icons/gi' import { IoSend } from 'react-icons/io5' import { HiHand } from 'react-icons/hi' -import { ChatMessageParams, ChatRole } from '../../../hooks/useOpenAI' interface SidebarInputProps { loading: boolean - submitMessage: (messages: ChatMessageParams[]) => void + submitMessage: (prompt: string) => void clearMessages: () => void chatIsEmpty: boolean cancelRequest: () => void @@ -33,7 +32,7 @@ export function SidebarInput({ }, [loading]) const handleSubmit = () => { - submitMessage([{ content: text, role: ChatRole.USER }]) + submitMessage(text) setText('') } diff --git a/src/components/Sidebar/chat/index.tsx b/src/components/Sidebar/chat/index.tsx index a2cdcd7e..a205e15e 100644 --- a/src/components/Sidebar/chat/index.tsx +++ b/src/components/Sidebar/chat/index.tsx @@ -1,21 +1,23 @@ -import React, { useEffect } from 'react' +import { useEffect } from 'react' import ChatList from './chatList' import { SidebarInput } from './chatInput' -import { useChatCompletion } from '../../../hooks/useOpenAI' +import { useChatCompletion } from '../../../hooks/useChatCompletion' import { SYSTEM_PROMPT } from '../../../config/prompts' import { Settings } from '../../../config/settings' interface ChatProps { settings: Settings + chatId: string } -const Chat = ({ settings }: ChatProps) => { - const { messages, submitQuery, clearMessages, loading, cancelRequest } = +const Chat = ({ settings, chatId }: ChatProps) => { + const { messages, submitQuery, clearMessages, generating, cancelRequest } = useChatCompletion({ model: settings.chat.modal, apiKey: settings.chat.openAIKey!, mode: settings.chat.mode, systemPrompt: SYSTEM_PROMPT, + chatId }) useEffect(() => { @@ -25,7 +27,7 @@ const Chat = ({ settings }: ChatProps) => { prompt: string } if (action === 'generate') { - submitQuery([{ content: prompt, role: 'user' }]) + submitQuery(prompt) } } window.addEventListener('message', handleWindowMessage) @@ -39,7 +41,7 @@ const Chat = ({ settings }: ChatProps) => { <> { - const chat = new ChatOpenAI({ + const { messages, updateAssistantMessage, updateStoredMessages, clearMessages } = useCurrentChat(chatId) + const [generating, setGenerating] = useState(false) + + const controller = new AbortController(); + + const chat = useMemo(() => new ChatOpenAI({ streaming: true, - }); + }), []) + + + + const submitQuery = async (query: string) => { + const previousMessages = messages.map((msg) => { + switch (msg.role) { + case Role.ASSISTANT: + return new AIMessage(msg.message) + case Role.SYSTEM: + return new SystemMessage(msg.message) + case Role.USER: + return new HumanMessage(msg.message) + } + }) + setGenerating(true) + const response = await chat.call([...previousMessages, new HumanMessage(query)], { + signal: controller.signal, + callbacks: [ + { + handleLLMNewToken: updateAssistantMessage + }, + ], + }); + setGenerating(false) + updateStoredMessages() + return response.content + } + + const cancelRequest = () => { + controller.abort() + setGenerating(false) + } - const [messages, setMessages] = useStorage() + return { messages, submitQuery, generating, cancelRequest, clearMessages } } diff --git a/src/hooks/useCurrentChat.tsx b/src/hooks/useCurrentChat.tsx new file mode 100644 index 00000000..014c959a --- /dev/null +++ b/src/hooks/useCurrentChat.tsx @@ -0,0 +1,55 @@ +import { useEffect } from "react" +import { useState } from "react" +import { useStorage } from "./useStorage" + + +export enum Role { + "USER", + "ASSISTANT", + "SYSTEM" +} + +export type Message = { + role: Role + message: string + timestamp: number +} + +export const useCurrentChat = (chatId: string) => { + const [storedMessages, setStoredMessages] = useStorage(`CHAT-${chatId}`, []) + const [messages, setMessages] = useState(storedMessages) // we don't directly update storedMessages for performance reasons + + const updateAssistantMessage = (chunk: string) => { + setMessages(messages => { + const lastMessage = messages[messages.length - 1] + lastMessage.message += chunk + return [...messages] + }) + } + + const addNewMessage = (role: Role, message: string) => { + const newMessage: Message = { + role, + message, + timestamp: Date.now(), + } + setMessages([...messages, newMessage]) + } + + const updateStoredMessages = () => { + setStoredMessages(messages) + } + + const clearMessages = () => { + setMessages([]) + updateStoredMessages() + } + + return { + messages, + updateAssistantMessage, + addNewMessage, + updateStoredMessages, + clearMessages + } +} diff --git a/src/hooks/useHistory.tsx b/src/hooks/useHistory.tsx index 436d6a9a..91a2de1f 100644 --- a/src/hooks/useHistory.tsx +++ b/src/hooks/useHistory.tsx @@ -11,12 +11,12 @@ interface ChatHistory { export const useHistory = () => { const [history, setHistory] = useStorage("HISTORY", []) - const createChatHistory = () => { + const createChatHistory = (name: string) => { const newId = randomUUID() setHistory(prev => [...prev, { id: newId, - name: "", + name, createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(), }])