From e8ed4df31aaac4001192f64a66701c48bc6e9378 Mon Sep 17 00:00:00 2001 From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com> Date: Wed, 28 Aug 2024 14:27:22 +0530 Subject: [PATCH] feat(chat-window): close socket on unmount --- ui/components/ChatWindow.tsx | 410 ++++++++++++++++++----------------- 1 file changed, 208 insertions(+), 202 deletions(-) diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx index cc93da8..9a1fe3c 100644 --- a/ui/components/ChatWindow.tsx +++ b/ui/components/ChatWindow.tsx @@ -1,42 +1,42 @@ -'use client'; +'use client' -import { useEffect, useRef, useState } from 'react'; -import { Document } from '@langchain/core/documents'; -import Navbar from './Navbar'; -import Chat from './Chat'; -import EmptyChat from './EmptyChat'; -import crypto from 'crypto'; -import { toast } from 'sonner'; -import { useSearchParams } from 'next/navigation'; -import { getSuggestions } from '@/lib/actions'; -import Error from 'next/error'; +import { useEffect, useRef, useState } from 'react' +import { Document } from '@langchain/core/documents' +import Navbar from './Navbar' +import Chat from './Chat' +import EmptyChat from './EmptyChat' +import crypto from 'crypto' +import { toast } from 'sonner' +import { useSearchParams } from 'next/navigation' +import { getSuggestions } from '@/lib/actions' +import Error from 'next/error' export type Message = { - messageId: string; - chatId: string; - createdAt: Date; - content: string; - role: 'user' | 'assistant'; - suggestions?: string[]; - sources?: Document[]; -}; + messageId: string + chatId: string + createdAt: Date + content: string + role: 'user' | 'assistant' + suggestions?: string[] + sources?: Document[] +} const useSocket = ( url: string, setIsWSReady: (ready: boolean) => void, setError: (error: boolean) => void, ) => { - const [ws, setWs] = useState(null); + const [ws, setWs] = useState(null) useEffect(() => { if (!ws) { const connectWs = async () => { - let chatModel = localStorage.getItem('chatModel'); - let chatModelProvider = localStorage.getItem('chatModelProvider'); - let embeddingModel = localStorage.getItem('embeddingModel'); + let chatModel = localStorage.getItem('chatModel') + let chatModelProvider = localStorage.getItem('chatModelProvider') + let embeddingModel = localStorage.getItem('embeddingModel') let embeddingModelProvider = localStorage.getItem( 'embeddingModelProvider', - ); + ) const providers = await fetch( `${process.env.NEXT_PUBLIC_API_URL}/models`, @@ -45,7 +45,7 @@ const useSocket = ( 'Content-Type': 'application/json', }, }, - ).then(async (res) => await res.json()); + ).then(async res => await res.json()) if ( !chatModel || @@ -54,58 +54,55 @@ const useSocket = ( !embeddingModelProvider ) { if (!chatModel || !chatModelProvider) { - const chatModelProviders = providers.chatModelProviders; + const chatModelProviders = providers.chatModelProviders - chatModelProvider = Object.keys(chatModelProviders)[0]; + chatModelProvider = Object.keys(chatModelProviders)[0] if (chatModelProvider === 'custom_openai') { toast.error( 'Seems like you are using the custom OpenAI provider, please open the settings and configure the API key and base URL', - ); - setError(true); - return; + ) + setError(true) + return } else { - chatModel = Object.keys(chatModelProviders[chatModelProvider])[0]; + chatModel = Object.keys(chatModelProviders[chatModelProvider])[0] if ( !chatModelProviders || Object.keys(chatModelProviders).length === 0 ) - return toast.error('No chat models available'); + return toast.error('No chat models available') } } if (!embeddingModel || !embeddingModelProvider) { - const embeddingModelProviders = providers.embeddingModelProviders; + const embeddingModelProviders = providers.embeddingModelProviders if ( !embeddingModelProviders || Object.keys(embeddingModelProviders).length === 0 ) - return toast.error('No embedding models available'); + return toast.error('No embedding models available') - embeddingModelProvider = Object.keys(embeddingModelProviders)[0]; + embeddingModelProvider = Object.keys(embeddingModelProviders)[0] embeddingModel = Object.keys( embeddingModelProviders[embeddingModelProvider], - )[0]; + )[0] } - localStorage.setItem('chatModel', chatModel!); - localStorage.setItem('chatModelProvider', chatModelProvider); - localStorage.setItem('embeddingModel', embeddingModel!); - localStorage.setItem( - 'embeddingModelProvider', - embeddingModelProvider, - ); + localStorage.setItem('chatModel', chatModel!) + localStorage.setItem('chatModelProvider', chatModelProvider) + localStorage.setItem('embeddingModel', embeddingModel!) + localStorage.setItem('embeddingModelProvider', embeddingModelProvider) } else { - const chatModelProviders = providers.chatModelProviders; - const embeddingModelProviders = providers.embeddingModelProviders; + const chatModelProviders = providers.chatModelProviders + const embeddingModelProviders = providers.embeddingModelProviders if ( Object.keys(chatModelProviders).length > 0 && !chatModelProviders[chatModelProvider] ) { - chatModelProvider = Object.keys(chatModelProviders)[0]; - localStorage.setItem('chatModelProvider', chatModelProvider); + chatModelProvider = Object.keys(chatModelProviders)[0] + localStorage.setItem('chatModelProvider', chatModelProvider) } if ( @@ -113,19 +110,19 @@ const useSocket = ( chatModelProvider != 'custom_openai' && !chatModelProviders[chatModelProvider][chatModel] ) { - chatModel = Object.keys(chatModelProviders[chatModelProvider])[0]; - localStorage.setItem('chatModel', chatModel); + chatModel = Object.keys(chatModelProviders[chatModelProvider])[0] + localStorage.setItem('chatModel', chatModel) } if ( Object.keys(embeddingModelProviders).length > 0 && !embeddingModelProviders[embeddingModelProvider] ) { - embeddingModelProvider = Object.keys(embeddingModelProviders)[0]; + embeddingModelProvider = Object.keys(embeddingModelProviders)[0] localStorage.setItem( 'embeddingModelProvider', embeddingModelProvider, - ); + ) } if ( @@ -134,77 +131,77 @@ const useSocket = ( ) { embeddingModel = Object.keys( embeddingModelProviders[embeddingModelProvider], - )[0]; - localStorage.setItem('embeddingModel', embeddingModel); + )[0] + localStorage.setItem('embeddingModel', embeddingModel) } } - const wsURL = new URL(url); - const searchParams = new URLSearchParams({}); + const wsURL = new URL(url) + const searchParams = new URLSearchParams({}) - searchParams.append('chatModel', chatModel!); - searchParams.append('chatModelProvider', chatModelProvider); + searchParams.append('chatModel', chatModel!) + searchParams.append('chatModelProvider', chatModelProvider) if (chatModelProvider === 'custom_openai') { searchParams.append( 'openAIApiKey', localStorage.getItem('openAIApiKey')!, - ); + ) searchParams.append( 'openAIBaseURL', localStorage.getItem('openAIBaseURL')!, - ); + ) } - searchParams.append('embeddingModel', embeddingModel!); - searchParams.append('embeddingModelProvider', embeddingModelProvider); + searchParams.append('embeddingModel', embeddingModel!) + searchParams.append('embeddingModelProvider', embeddingModelProvider) - wsURL.search = searchParams.toString(); + wsURL.search = searchParams.toString() - const ws = new WebSocket(wsURL.toString()); + const ws = new WebSocket(wsURL.toString()) const timeoutId = setTimeout(() => { if (ws.readyState !== 1) { toast.error( 'Failed to connect to the server. Please try again later.', - ); + ) } - }, 10000); + }, 10000) ws.onopen = () => { - console.log('[DEBUG] open'); - clearTimeout(timeoutId); - setIsWSReady(true); - }; + console.log('[DEBUG] open') + clearTimeout(timeoutId) + setIsWSReady(true) + } ws.onerror = () => { - clearTimeout(timeoutId); - setError(true); - toast.error('WebSocket connection error.'); - }; + clearTimeout(timeoutId) + setError(true) + toast.error('WebSocket connection error.') + } ws.onclose = () => { - clearTimeout(timeoutId); - setError(true); - console.log('[DEBUG] closed'); - }; + clearTimeout(timeoutId) + setError(true) + console.log('[DEBUG] closed') + } - ws.addEventListener('message', (e) => { - const data = JSON.parse(e.data); + ws.addEventListener('message', e => { + const data = JSON.parse(e.data) if (data.type === 'error') { - toast.error(data.data); + toast.error(data.data) } - }); + }) - setWs(ws); - }; + setWs(ws) + } - connectWs(); + connectWs() } - }, [ws, url, setIsWSReady, setError]); + }, [ws, url, setIsWSReady, setError]) - return ws; -}; + return ws +} const loadMessages = async ( chatId: string, @@ -222,66 +219,66 @@ const loadMessages = async ( 'Content-Type': 'application/json', }, }, - ); + ) if (res.status === 404) { - setNotFound(true); - setIsMessagesLoaded(true); - return; + setNotFound(true) + setIsMessagesLoaded(true) + return } - const data = await res.json(); + const data = await res.json() const messages = data.messages.map((msg: any) => { return { ...msg, ...JSON.parse(msg.metadata), - }; - }) as Message[]; + } + }) as Message[] - setMessages(messages); + setMessages(messages) - const history = messages.map((msg) => { - return [msg.role, msg.content]; - }) as [string, string][]; + const history = messages.map(msg => { + return [msg.role, msg.content] + }) as [string, string][] - console.log('[DEBUG] messages loaded'); + console.log('[DEBUG] messages loaded') - document.title = messages[0].content; + document.title = messages[0].content - setChatHistory(history); - setFocusMode(data.chat.focusMode); - setIsMessagesLoaded(true); -}; + setChatHistory(history) + setFocusMode(data.chat.focusMode) + setIsMessagesLoaded(true) +} const ChatWindow = ({ id }: { id?: string }) => { - const searchParams = useSearchParams(); - const initialMessage = searchParams.get('q'); + const searchParams = useSearchParams() + const initialMessage = searchParams.get('q') - const [chatId, setChatId] = useState(id); - const [newChatCreated, setNewChatCreated] = useState(false); + const [chatId, setChatId] = useState(id) + const [newChatCreated, setNewChatCreated] = useState(false) - const [hasError, setHasError] = useState(false); - const [isReady, setIsReady] = useState(false); + const [hasError, setHasError] = useState(false) + const [isReady, setIsReady] = useState(false) - const [isWSReady, setIsWSReady] = useState(false); + const [isWSReady, setIsWSReady] = useState(false) const ws = useSocket( process.env.NEXT_PUBLIC_WS_URL!, setIsWSReady, setHasError, - ); + ) - const [loading, setLoading] = useState(false); - const [messageAppeared, setMessageAppeared] = useState(false); + const [loading, setLoading] = useState(false) + const [messageAppeared, setMessageAppeared] = useState(false) - const [chatHistory, setChatHistory] = useState<[string, string][]>([]); - const [messages, setMessages] = useState([]); + const [chatHistory, setChatHistory] = useState<[string, string][]>([]) + const [messages, setMessages] = useState([]) - const [focusMode, setFocusMode] = useState('webSearch'); + const [focusMode, setFocusMode] = useState('webSearch') - const [isMessagesLoaded, setIsMessagesLoaded] = useState(false); + const [isMessagesLoaded, setIsMessagesLoaded] = useState(false) - const [notFound, setNotFound] = useState(false); + const [notFound, setNotFound] = useState(false) useEffect(() => { if ( @@ -297,37 +294,46 @@ const ChatWindow = ({ id }: { id?: string }) => { setChatHistory, setFocusMode, setNotFound, - ); + ) } else if (!chatId) { - setNewChatCreated(true); - setIsMessagesLoaded(true); - setChatId(crypto.randomBytes(20).toString('hex')); + setNewChatCreated(true) + setIsMessagesLoaded(true) + setChatId(crypto.randomBytes(20).toString('hex')) } // eslint-disable-next-line react-hooks/exhaustive-deps - }, []); - - const messagesRef = useRef([]); + }, []) useEffect(() => { - messagesRef.current = messages; - }, [messages]); + return () => { + if (ws?.readyState === 1) { + ws.close() + console.log('[DEBUG] closed') + } + } + }, []) + + const messagesRef = useRef([]) + + useEffect(() => { + messagesRef.current = messages + }, [messages]) useEffect(() => { if (isMessagesLoaded && isWSReady) { - setIsReady(true); + setIsReady(true) } - }, [isMessagesLoaded, isWSReady]); + }, [isMessagesLoaded, isWSReady]) const sendMessage = async (message: string) => { - if (loading) return; - setLoading(true); - setMessageAppeared(false); + if (loading) return + setLoading(true) + setMessageAppeared(false) - let sources: Document[] | undefined = undefined; - let recievedMessage = ''; - let added = false; + let sources: Document[] | undefined = undefined + let recievedMessage = '' + let added = false - const messageId = crypto.randomBytes(7).toString('hex'); + const messageId = crypto.randomBytes(7).toString('hex') ws?.send( JSON.stringify({ @@ -339,9 +345,9 @@ const ChatWindow = ({ id }: { id?: string }) => { focusMode: focusMode, history: [...chatHistory, ['human', message]], }), - ); + ) - setMessages((prevMessages) => [ + setMessages(prevMessages => [ ...prevMessages, { content: message, @@ -350,21 +356,21 @@ const ChatWindow = ({ id }: { id?: string }) => { role: 'user', createdAt: new Date(), }, - ]); + ]) const messageHandler = async (e: MessageEvent) => { - const data = JSON.parse(e.data); + const data = JSON.parse(e.data) if (data.type === 'error') { - toast.error(data.data); - setLoading(false); - return; + toast.error(data.data) + setLoading(false) + return } if (data.type === 'sources') { - sources = data.data; + sources = data.data if (!added) { - setMessages((prevMessages) => [ + setMessages(prevMessages => [ ...prevMessages, { content: '', @@ -374,15 +380,15 @@ const ChatWindow = ({ id }: { id?: string }) => { sources: sources, createdAt: new Date(), }, - ]); - added = true; + ]) + added = true } - setMessageAppeared(true); + setMessageAppeared(true) } if (data.type === 'message') { if (!added) { - setMessages((prevMessages) => [ + setMessages(prevMessages => [ ...prevMessages, { content: data.data, @@ -392,35 +398,35 @@ const ChatWindow = ({ id }: { id?: string }) => { sources: sources, createdAt: new Date(), }, - ]); - added = true; + ]) + added = true } - setMessages((prev) => - prev.map((message) => { + setMessages(prev => + prev.map(message => { if (message.messageId === data.messageId) { - return { ...message, content: message.content + data.data }; + return { ...message, content: message.content + data.data } } - return message; + return message }), - ); + ) - recievedMessage += data.data; - setMessageAppeared(true); + recievedMessage += data.data + setMessageAppeared(true) } if (data.type === 'messageEnd') { - setChatHistory((prevHistory) => [ + setChatHistory(prevHistory => [ ...prevHistory, ['human', message], ['assistant', recievedMessage], - ]); + ]) - ws?.removeEventListener('message', messageHandler); - setLoading(false); + ws?.removeEventListener('message', messageHandler) + setLoading(false) - const lastMsg = messagesRef.current[messagesRef.current.length - 1]; + const lastMsg = messagesRef.current[messagesRef.current.length - 1] if ( lastMsg.role === 'assistant' && @@ -428,54 +434,54 @@ const ChatWindow = ({ id }: { id?: string }) => { lastMsg.sources.length > 0 && !lastMsg.suggestions ) { - const suggestions = await getSuggestions(messagesRef.current); - setMessages((prev) => - prev.map((msg) => { + const suggestions = await getSuggestions(messagesRef.current) + setMessages(prev => + prev.map(msg => { if (msg.messageId === lastMsg.messageId) { - return { ...msg, suggestions: suggestions }; + return { ...msg, suggestions: suggestions } } - return msg; + return msg }), - ); + ) } } - }; + } - ws?.addEventListener('message', messageHandler); - }; + ws?.addEventListener('message', messageHandler) + } const rewrite = (messageId: string) => { - const index = messages.findIndex((msg) => msg.messageId === messageId); + const index = messages.findIndex(msg => msg.messageId === messageId) - if (index === -1) return; + if (index === -1) return - const message = messages[index - 1]; + const message = messages[index - 1] - setMessages((prev) => { - return [...prev.slice(0, messages.length > 2 ? index - 1 : 0)]; - }); - setChatHistory((prev) => { - return [...prev.slice(0, messages.length > 2 ? index - 1 : 0)]; - }); + setMessages(prev => { + return [...prev.slice(0, messages.length > 2 ? index - 1 : 0)] + }) + setChatHistory(prev => { + return [...prev.slice(0, messages.length > 2 ? index - 1 : 0)] + }) - sendMessage(message.content); - }; + sendMessage(message.content) + } useEffect(() => { if (isReady && initialMessage) { - sendMessage(initialMessage); + sendMessage(initialMessage) } // eslint-disable-next-line react-hooks/exhaustive-deps - }, [isReady, initialMessage]); + }, [isReady, initialMessage]) if (hasError) { return ( -
-

+

+

Failed to connect to the server. Please try again later.

- ); + ) } return isReady ? ( @@ -504,25 +510,25 @@ const ChatWindow = ({ id }: { id?: string }) => {
) ) : ( -
+
- ); -}; + ) +} -export default ChatWindow; +export default ChatWindow