From 7f9651a3053cb94fe51284eb9e6fb152036a42ca Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Fri, 5 May 2023 18:49:00 -0300 Subject: [PATCH] fixed onClose method websocket, streaming for text working --- src/frontend/src/modals/chatModal/index.tsx | 192 +++++++++++--------- 1 file changed, 108 insertions(+), 84 deletions(-) diff --git a/src/frontend/src/modals/chatModal/index.tsx b/src/frontend/src/modals/chatModal/index.tsx index 8b94fb141..40f0e88a2 100644 --- a/src/frontend/src/modals/chatModal/index.tsx +++ b/src/frontend/src/modals/chatModal/index.tsx @@ -26,9 +26,15 @@ export default function ChatModal({ const [chatHistory, setChatHistory] = useState([]); const { reactFlowInstance } = useContext(typesContext); const { setErrorData, setNoticeData } = useContext(alertContext); - const [isStreaming,setIsStreaming] = useState() - const [ws, setWs] = useState(null); + const ws = useRef(null) const [lockChat, setLockChat] = useState(false); + const isOpen = useRef(open); + + useEffect(() => { + isOpen.current = open; + }, [open]); + var isStream = false; + const addChatHistory = ( message: string, isSend: boolean, @@ -47,9 +53,86 @@ export default function ChatModal({ return newChat; }); }; + - function checkOpenFunction(){ - return open + function updateLastMessage(str: string) { + console.log(str); + setChatHistory((old) => { + let newChat = [...old]; + newChat[newChat.length - 1].message = + newChat[newChat.length - 1].message + str; + return newChat; + }); + } + + function handleOnClose(event: CloseEvent) { + if (isOpen.current) { + setLockChat(false); + setTimeout(() => { + connectWS() + }, 1000); + } + } + + function handleWsMessage(data: any) { + if (Array.isArray(data)) { + //set chat history + setChatHistory((_) => { + let newChatHistory: ChatMessageType[] = []; + data.forEach( + (chatItem: { + intermediate_steps?: "string"; + is_bot: boolean; + message: string; + type: string; + files?: Array; + }) => { + if (chatItem.message) { + newChatHistory.push( + chatItem.files + ? { + isSend: !chatItem.is_bot, + message: chatItem.message, + thought: chatItem.intermediate_steps, + files: chatItem.files, + } + : { + isSend: !chatItem.is_bot, + message: chatItem.message, + thought: chatItem.intermediate_steps, + } + ); + } + } + ); + return newChatHistory; + }); + } + if (data.type === "start") { + console.log("start"); + addChatHistory("", false); + isStream = true; + } + if (data.type === "end") { + setLockChat(false); + isStream = false; + // if (data.files) { + // addChatHistory( + // data.message, + // false, + // data.intermediate_steps, + // data.files + // ); + // } else { + // addChatHistory(data.message, false, data.intermediate_steps); + // } + } + if (data.type === "file") { + console.log(data); + } + if (data.type === "stream" && isStream) { + updateLastMessage(data.message); + } } function connectWS() { @@ -65,76 +148,28 @@ export default function ChatModal({ newWs.onopen = () => { console.log("WebSocket connection established!"); }; - console.log(flow.id) + console.log(flow.id); newWs.onmessage = (event) => { - setLockChat(false); const data = JSON.parse(event.data); console.log("Received data:", data); + handleWsMessage(data); //get chat history - if (Array.isArray(data)) { - console.log(data); - - setChatHistory((_) => { - let newChatHistory: ChatMessageType[] = []; - data.forEach( - (chatItem: { - intermediate_steps?: "string"; - is_bot: boolean; - message: string; - type: string; - files?: Array; - }) => { - if (chatItem.message) { - newChatHistory.push( - chatItem.files - ? { - isSend: !chatItem.is_bot, - message: chatItem.message, - thought: chatItem.intermediate_steps, - files: chatItem.files, - } - : { - isSend: !chatItem.is_bot, - message: chatItem.message, - thought: chatItem.intermediate_steps, - } - ); - } - } - ); - return newChatHistory; - }); - } - if (data.type === "end") { - if (data.files) { - addChatHistory( - data.message, - false, - data.intermediate_steps, - data.files - ); - } else { - addChatHistory(data.message, false, data.intermediate_steps); - } - } - if (data.type === "file") { - console.log(data); - } }; - newWs.onclose = (_) => { - if(checkOpenFunction){ - setLockChat(false); - setTimeout(() => { - connectWS(); - }, 1000); - } + newWs.onclose = (event) => { + handleOnClose(event); }; newWs.onerror = (ev) => { console.log(ev, "error"); + setErrorData({ + title: "There was an error on web connection, please: ", + list: [ + "Refresh the page", + "Use a new flow tab", + "Check if the backend is up", + ], + }); }; - setWs(newWs); - - return newWs; + ws.current=newWs } catch { setErrorData({ title: "There was an error on web connection, please: ", @@ -148,25 +183,21 @@ export default function ChatModal({ } useEffect(() => { - let newWs = connectWS(); + connectWS(); return () => { - newWs.close(); + console.log("unmount") + console.log(ws) + if (ws) { + ws.current.close(); + } }; }, []); - useEffect(() => { - if (ws && (ws.readyState === ws.CLOSED || ws.readyState === ws.CLOSING)) { - let newWs = connectWS(); - return () => { - newWs.close(); - }; - } - }, [lockChat]); async function sendAll(data: sendAllProps) { try { if (ws) { - ws.send(JSON.stringify(data)); + ws.current.send(JSON.stringify(data)); } } catch (error) { setErrorData({ @@ -182,12 +213,6 @@ export default function ChatModal({ if (ref.current) ref.current.scrollIntoView({ behavior: "smooth" }); }, [chatHistory]); - useEffect(() => { - if (ws && ws.readyState === ws.CLOSED) { - setLockChat(false); - } - }, [lockChat]); - function validateNode(n: NodeType): Array { if (!n.data?.node?.template || !Object.keys(n.data.node.template)) { setNoticeData({ @@ -244,7 +269,6 @@ export default function ChatModal({ let message = chatValue; setChatValue(""); addChatHistory(message, true); - sendAll({ ...reactFlowInstance.toObject(), message, @@ -267,7 +291,7 @@ export default function ChatModal({ } function clearChat() { setChatHistory([]); - ws.send(JSON.stringify({ clear_history: true })); + ws.current.send(JSON.stringify({ clear_history: true })); } function setModalOpen(x: boolean) {