fix: stops stream upon conversation

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;

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();
+  }
+};

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'));

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();

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;

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();
 }