From 23de688bf32f30f17b85f95dc2bd7ff0c6069f5d Mon Sep 17 00:00:00 2001 From: Daniel Avila Date: Sat, 11 Mar 2023 21:42:08 -0500 Subject: [PATCH] fix: stops stream upon conversation --- api/app/detectCode.js | 4 +- api/server/routes/stopStream.js | 11 + .../components/Conversations/Conversation.jsx | 15 +- client/src/components/Main/TextChat.jsx | 294 +++++++++++------- client/src/store/submitSlice.js | 8 +- client/src/utils/handleSubmit.js | 6 + 6 files changed, 218 insertions(+), 120 deletions(-) create mode 100644 api/server/routes/stopStream.js diff --git a/api/app/detectCode.js b/api/app/detectCode.js index b631d1921e..0d6fad0534 100644 --- a/api/app/detectCode.js +++ b/api/app/detectCode.js @@ -2,7 +2,7 @@ const { ModelOperations } = require('@vscode/vscode-languagedetection'); const languages = require('../utils/languages.js'); const codeRegex = /(```[\s\S]*?```)/g; // const languageMatch = /```(\w+)/; -const replaceRegex = /```\w+/g; +const replaceRegex = /```\w+\n/g; const detectCode = async (input) => { try { @@ -22,7 +22,7 @@ const detectCode = async (input) => { } console.log('[detectCode.js] replacing', match, 'with', '```shell'); - text = text.replace(match, '```shell'); + text = text.replace(match, '```shell\n'); }); return text; diff --git a/api/server/routes/stopStream.js b/api/server/routes/stopStream.js new file mode 100644 index 0000000000..f12ef34de2 --- /dev/null +++ b/api/server/routes/stopStream.js @@ -0,0 +1,11 @@ +module.exports = (req, res, next) => { + let { stopStream } = req.body; + if (stopStream) { + console.log('stopStream'); + res.write('event: stop\ndata:\n\n'); + res.end(); + return; + } else { + next(); + } +}; diff --git a/client/src/components/Conversations/Conversation.jsx b/client/src/components/Conversations/Conversation.jsx index 470206cb51..735c2f74e8 100644 --- a/client/src/components/Conversations/Conversation.jsx +++ b/client/src/components/Conversations/Conversation.jsx @@ -3,7 +3,7 @@ import RenameButton from './RenameButton'; import DeleteButton from './DeleteButton'; import { useSelector, useDispatch } from 'react-redux'; import { setConversation } from '~/store/convoSlice'; -import { setStopStream, setCustomGpt, setModel, setCustomModel } from '~/store/submitSlice'; +import { setSubmitState, setSubmission, setStopStream, setCustomGpt, setModel, setCustomModel } from '~/store/submitSlice'; import { setMessages, setEmptyMessage } from '~/store/messageSlice'; import { setText } from '~/store/textSlice'; import manualSWR from '~/utils/fetchers'; @@ -21,6 +21,7 @@ export default function Conversation({ const [renaming, setRenaming] = useState(false); const [titleInput, setTitleInput] = useState(title); const { modelMap } = useSelector((state) => state.models); + const { stopStream } = useSelector((state) => state.submit); const inputRef = useRef(null); const dispatch = useDispatch(); const { trigger, isMutating } = manualSWR(`/api/messages/${id}`, 'get'); @@ -31,7 +32,11 @@ export default function Conversation({ return; } - dispatch(setStopStream(true)); + if (!stopStream) { + dispatch(setStopStream(true)); + dispatch(setSubmission({})); + dispatch(setSubmitState(false)); + } dispatch(setEmptyMessage()); const convo = { title, error: false, conversationId: id, chatGptLabel, promptPrefix }; @@ -67,12 +72,6 @@ export default function Conversation({ ); } const data = await trigger(); - // while (isMutating) { - // await new Promise((resolve) => setTimeout(() => { - // dispatch(setMessages([])); - // resolve(); - // }, 50)); - // } if (chatGptLabel) { dispatch(setModel('chatgptCustom')); diff --git a/client/src/components/Main/TextChat.jsx b/client/src/components/Main/TextChat.jsx index 9fe119adef..41c9d35b5c 100644 --- a/client/src/components/Main/TextChat.jsx +++ b/client/src/components/Main/TextChat.jsx @@ -1,4 +1,5 @@ -import React, { useState } from 'react'; +import React, { useEffect, useState } from 'react'; +import { SSE } from '~/utils/sse'; import SubmitButton from './SubmitButton'; import Regenerate from './Regenerate'; import ModelMenu from '../Models/ModelMenu'; @@ -8,7 +9,7 @@ import handleSubmit from '~/utils/handleSubmit'; import { useSelector, useDispatch } from 'react-redux'; import { setConversation, setError } from '~/store/convoSlice'; import { setMessages } from '~/store/messageSlice'; -import { setSubmitState } from '~/store/submitSlice'; +import { setSubmitState, setSubmission } from '~/store/submitSlice'; import { setText } from '~/store/textSlice'; export default function TextChat({ messages }) { @@ -16,12 +17,98 @@ export default function TextChat({ messages }) { const dispatch = useDispatch(); const convo = useSelector((state) => state.convo); const { initial } = useSelector((state) => state.models); - const { isSubmitting, stopStream, disabled, model, chatGptLabel, promptPrefix } = useSelector( - (state) => state.submit - ); + const { isSubmitting, stopStream, submission, disabled, model, chatGptLabel, promptPrefix } = + useSelector((state) => state.submit); const { text } = useSelector((state) => state.text); const { error } = convo; - const isCustomModel = model === 'chatgptCustom' || !initial[model]; + + const messageHandler = (data, currentState) => { + const { messages, currentMsg, sender } = currentState; + dispatch(setMessages([...messages, currentMsg, { sender, text: data }])); + }; + + const convoHandler = (data, currentState) => { + const { messages, currentMsg, sender, isCustomModel, model, chatGptLabel, promptPrefix } = + currentState; + dispatch( + setMessages([...messages, currentMsg, { sender, text: data.text || data.response }]) + ); + + const isBing = model === 'bingai' || model === 'sydney'; + + if (!isBing && convo.conversationId === null && convo.parentMessageId === null) { + const { title, conversationId, id } = data; + dispatch( + setConversation({ + title, + conversationId, + parentMessageId: id, + jailbreakConversationId: null, + conversationSignature: null, + clientId: null, + invocationId: null, + chatGptLabel: model === isCustomModel ? chatGptLabel : null, + promptPrefix: model === isCustomModel ? promptPrefix : null + }) + ); + } else if ( + model === 'bingai' && + convo.conversationId === null && + convo.invocationId === null + ) { + console.log('Bing data:', data); + const { title, conversationSignature, clientId, conversationId, invocationId } = data; + dispatch( + setConversation({ + title, + parentMessageId: null, + conversationSignature, + clientId, + conversationId, + invocationId + }) + ); + } else if (model === 'sydney') { + const { + title, + jailbreakConversationId, + parentMessageId, + conversationSignature, + clientId, + conversationId, + invocationId + } = data; + dispatch( + setConversation({ + title, + jailbreakConversationId, + parentMessageId, + conversationSignature, + clientId, + conversationId, + invocationId + }) + ); + } + + dispatch(setSubmitState(false)); + }; + + const errorHandler = (event, currentState) => { + const { initialResponse, messages, currentMsg, message } = currentState; + console.log('Error:', event); + const errorResponse = { + ...initialResponse, + text: `An error occurred. Please try again in a few moments.\n\nError message: ${event.data}`, + error: true + }; + setErrorMessage(event.data); + dispatch(setSubmitState(false)); + dispatch(setMessages([...messages.slice(0, -2), currentMsg, errorResponse])); + dispatch(setText(message)); + dispatch(setError(true)); + return; + }; const submitMessage = () => { if (error) { @@ -31,124 +118,115 @@ export default function TextChat({ messages }) { if (!!isSubmitting || text.trim() === '') { return; } - dispatch(setSubmitState(true)); + + const isCustomModel = model === 'chatgptCustom' || !initial[model]; const message = text.trim(); const currentMsg = { sender: 'User', text: message, current: true }; const sender = model === 'chatgptCustom' ? chatGptLabel : model; const initialResponse = { sender, text: '' }; + + dispatch(setSubmitState(true)); dispatch(setMessages([...messages, currentMsg, initialResponse])); dispatch(setText('')); - const messageHandler = (data, events) => { - if (stopStream) { - console.log('Stopping stream'); - events.close(); - return; - } - dispatch(setMessages([...messages, currentMsg, { sender, text: data }])); - }; - const convoHandler = (data) => { - dispatch( - setMessages([...messages, currentMsg, { sender, text: data.text || data.response }]) - ); - const isBing = model === 'bingai' || model === 'sydney'; - - if ( - !isBing && - convo.conversationId === null && - convo.parentMessageId === null - ) { - const { title, conversationId, id } = data; - dispatch( - setConversation({ - title, - conversationId, - parentMessageId: id, - jailbreakConversationId: null, - conversationSignature: null, - clientId: null, - invocationId: null, - chatGptLabel: model === isCustomModel ? chatGptLabel : null, - promptPrefix: model === isCustomModel ? promptPrefix : null - }) - ); - } else if ( - model === 'bingai' && - convo.conversationId === null && - convo.invocationId === null - ) { - console.log('Bing data:', data) - const { - title, - conversationSignature, - clientId, - conversationId, - invocationId - } = data; - dispatch( - setConversation({ - title, - parentMessageId: null, - conversationSignature, - clientId, - conversationId, - invocationId, - }) - ); - } else if (model === 'sydney') { - const { - title, - jailbreakConversationId, - parentMessageId, - conversationSignature, - clientId, - conversationId, - invocationId - } = data; - dispatch( - setConversation({ - title, - jailbreakConversationId, - parentMessageId, - conversationSignature, - clientId, - conversationId, - invocationId, - }) - ); - } - - dispatch(setSubmitState(false)); - }; - - const errorHandler = (event) => { - console.log('Error:', event); - const errorResponse = { - ...initialResponse, - text: `An error occurred. Please try again in a few moments.\n\nError message: ${event.data}`, - error: true - }; - setErrorMessage(event.data); - dispatch(setSubmitState(false)); - dispatch(setMessages([...messages.slice(0, -2), currentMsg, errorResponse])); - dispatch(setText(message)); - dispatch(setError(true)); - return; - }; const submission = { model, text: message, convo, - messageHandler, - convoHandler, - errorHandler, chatGptLabel, - promptPrefix + promptPrefix, + isCustomModel, + message, + messages, + currentMsg, + sender, + initialResponse }; console.log('User Input:', message); - handleSubmit(submission); + // handleSubmit(submission); + dispatch(setSubmission(submission)); }; + const createPayload = ({ model, text, convo, chatGptLabel, promptPrefix }) => { + const endpoint = `/api/ask`; + let payload = { model, text, chatGptLabel, promptPrefix }; + if (convo.conversationId && convo.parentMessageId) { + payload = { + ...payload, + conversationId: convo.conversationId, + parentMessageId: convo.parentMessageId + }; + } + + const isBing = model === 'bingai' || model === 'sydney'; + if (isBing && convo.conversationId) { + payload = { + ...payload, + jailbreakConversationId: convo.jailbreakConversationId, + conversationId: convo.conversationId, + conversationSignature: convo.conversationSignature, + clientId: convo.clientId, + invocationId: convo.invocationId + }; + } + + let server = endpoint; + server = model === 'bingai' ? server + '/bing' : server; + server = model === 'sydney' ? server + '/sydney' : server; + return { server, payload }; + }; + + useEffect(() => { + if (Object.keys(submission).length === 0) { + return; + } + + const currentState = submission; + const { server, payload } = createPayload(submission); + const onMessage = (e) => { + if (stopStream) { + return; + } + + const data = JSON.parse(e.data); + let text = data.text || data.response; + if (data.message) { + messageHandler(text, currentState); + } + + if (data.final) { + convoHandler(data, currentState); + console.log('final', data); + } else { + // console.log('dataStream', data); + } + }; + + const events = new SSE(server, { + payload: JSON.stringify(payload), + headers: { 'Content-Type': 'application/json' } + }); + + events.onopen = function () { + console.log('connection is opened'); + }; + + events.onmessage = onMessage; + + events.onerror = function (e) { + console.log('error in opening conn.'); + events.close(); + errorHandler(e, currentState); + }; + + events.stream(); + + return () => { + events.removeEventListener('message', onMessage); + events.close(); + }; + }, [submission]); + const handleKeyDown = (e) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); diff --git a/client/src/store/submitSlice.js b/client/src/store/submitSlice.js index cdcc230224..2a59f72e90 100644 --- a/client/src/store/submitSlice.js +++ b/client/src/store/submitSlice.js @@ -2,12 +2,13 @@ import { createSlice } from '@reduxjs/toolkit'; const initialState = { isSubmitting: false, + submission: {}, stopStream: false, disabled: false, model: 'chatgpt', promptPrefix: '', chatGptLabel: '', - customModel: null + customModel: null, }; const currentSlice = createSlice({ @@ -17,6 +18,9 @@ const currentSlice = createSlice({ setSubmitState: (state, action) => { state.isSubmitting = action.payload; }, + setSubmission: (state, action) => { + state.submission = action.payload; + }, setStopStream: (state, action) => { state.stopStream = action.payload; }, @@ -36,7 +40,7 @@ const currentSlice = createSlice({ } }); -export const { setSubmitState, setStopStream, setDisabled, setModel, setCustomGpt, setCustomModel } = +export const { setSubmitState, setSubmission, setStopStream, setDisabled, setModel, setCustomGpt, setCustomModel } = currentSlice.actions; export default currentSlice.reducer; diff --git a/client/src/utils/handleSubmit.js b/client/src/utils/handleSubmit.js index 1f6dfca027..5c4184aac5 100644 --- a/client/src/utils/handleSubmit.js +++ b/client/src/utils/handleSubmit.js @@ -68,5 +68,11 @@ export default function handleSubmit({ errorHandler(e); }; + events.addEventListener('stop', () => { + // Close the SSE stream + console.log('stop event received'); + events.close(); + }); + events.stream(); } \ No newline at end of file