🌿 feat: Multi-response Streaming (#3191)

* chore: comment back handlePlusCommand

* chore: ignore .git dir

* refactor: pass newConversation to `useSelectMention`

refactor: pass newConversation to Mention component

refactor: useChatFunctions for modular use of `ask` and `regenerate`

refactor: set latest message only for the first index in useChatFunctions

refactor: pass setLatestMessage to useChatFunctions

refactor: Pass setSubmission to useChatFunctions for submission handling

refactor: consolidate event handlers to separate hook from useSSE

WIP: additional response handlers

feat: responsive added convo, clears on new chat/navigating to chat, assistants excluded

feat: Add conversationByKeySelector to select any conversation by index

WIP: handle second submission with messages paired to root

* style: surface-primary-contrast

* refactor: remove unnecessary console.log statement in useChatFunctions

* refactor: Consolidate imports in ChatForm and Input hooks

* refactor: compositional usage of useSSE for multiple streams

* WIP: set latest 'multi' message

* WIP: first pass, added response streaming

* pass: performant multi-message stream

* fix: styling and message render

* second pass: modular, performant multi-stream

* fix: align parentMessageId of multiMessage

* refactor: move resetting latestMultiMessage

* chore: update footer text in Chat component

* fix: stop button styling

* fix: handle abortMessage request for multi-response

* clear messages but bug with latest message reset present

* fix: add delay for additional message generation

* fix: access LAST_CONVO_SETUP by index

* style: add div to prevent layout shift before hover buttons render

* chore: Update Message component styling for card messages

* chore: move hook use order

* fix: abort middleware using unsent field from req.body

* feat: support multi-response stream from initial message

* refactor: buildTree function to improve readability and remove unused code

* feat: add logger for frontend dev

* refactor: use depth to track if message is really last in its branch

* fix(buildTree): default export

* fix: share parent message Id and avoid duplication error for multi-response streams

* fix: prevent addedConvo reset to response convo

* feat: allow setting multi message as latest message to control which to respond to

* chore: wrap setSiblingIdxRev with useCallback

* chore: styling and allow editing messages

* style: styling fixes

* feat: Add "AddMultiConvo" component to Chat Header

* feat: prevent clearing added convos on endpoint, preset, mention, or modelSpec switch

* fix: message styling fixes, mainly related to code blocks

* fix: stop button visibility logic

* fix: Handle edge case in abortMiddleware for non-existant `abortControllers`

* refactor: optimize/memoize icons

* chore(GoogleClient): change info to debug logs

* style: active message styling

* style: prevent layout shift due to placeholder row

* chore: remove unused code

* fix: Update BaseClient to handle optional request body properties

* fix(ci): `onStart` now accepts 2 args, the 2nd being responseMessageId

* chore: bump data-provider

api/server/controllers/AskController.js

@@ -81,7 +81,7 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       promptTokens,
     });
 
-    const { abortController, onStart } = createAbortController(req, res, getAbortData);
+    const { abortController, onStart } = createAbortController(req, res, getAbortData, getReqData);
 
     res.on('close', () => {
       logger.debug('[AskController] Request closed');
@@ -144,7 +144,9 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       await saveMessage({ ...response, user });
     }
 
-    await saveMessage(userMessage);
+    if (!client.skipSaveUserMessage) {
+      await saveMessage(userMessage);
+    }
 
     if (addTitle && parentMessageId === Constants.NO_PARENT && newConvo) {
       addTitle(req, {

api/server/controllers/EditController.js

@@ -81,7 +81,7 @@ const EditController = async (req, res, next, initializeClient) => {
     promptTokens,
   });
 
-  const { abortController, onStart } = createAbortController(req, res, getAbortData);
+  const { abortController, onStart } = createAbortController(req, res, getAbortData, getReqData);
 
   res.on('close', () => {
     logger.debug('[EditController] Request closed');

api/server/middleware/abortMiddleware.js

@@ -9,23 +9,28 @@ const { abortRun } = require('./abortRun');
 const { logger } = require('~/config');
 
 async function abortMessage(req, res) {
-  let { abortKey, conversationId, endpoint } = req.body;
-
-  if (!abortKey && conversationId) {
-    abortKey = conversationId;
-  }
+  let { abortKey, endpoint } = req.body;
 
   if (isAssistantsEndpoint(endpoint)) {
     return await abortRun(req, res);
   }
 
+  const conversationId = abortKey?.split(':')?.[0] ?? req.user.id;
+
+  if (!abortControllers.has(abortKey) && abortControllers.has(conversationId)) {
+    abortKey = conversationId;
+  }
+
   if (!abortControllers.has(abortKey) && !res.headersSent) {
     return res.status(204).send({ message: 'Request not found' });
   }
 
-  const { abortController } = abortControllers.get(abortKey);
+  const { abortController } = abortControllers.get(abortKey) ?? {};
+  if (!abortController) {
+    return res.status(204).send({ message: 'Request not found' });
+  }
   const finalEvent = await abortController.abortCompletion();
-  logger.debug('[abortMessage] Aborted request', { abortKey });
+  logger.info('[abortMessage] Aborted request', { abortKey });
   abortControllers.delete(abortKey);
 
   if (res.headersSent && finalEvent) {
@@ -50,12 +55,32 @@ const handleAbort = () => {
   };
 };
 
-const createAbortController = (req, res, getAbortData) => {
+const createAbortController = (req, res, getAbortData, getReqData) => {
   const abortController = new AbortController();
   const { endpointOption } = req.body;
-  const onStart = (userMessage) => {
+
+  abortController.getAbortData = function () {
+    return getAbortData();
+  };
+
+  /**
+   * @param {TMessage} userMessage
+   * @param {string} responseMessageId
+   */
+  const onStart = (userMessage, responseMessageId) => {
     sendMessage(res, { message: userMessage, created: true });
     const abortKey = userMessage?.conversationId ?? req.user.id;
+    const prevRequest = abortControllers.get(abortKey);
+    if (prevRequest && prevRequest?.abortController) {
+      const data = prevRequest.abortController.getAbortData();
+      getReqData({ userMessage: data?.userMessage });
+      const addedAbortKey = `${abortKey}:${responseMessageId}`;
+      abortControllers.set(addedAbortKey, { abortController, ...endpointOption });
+      res.on('finish', function () {
+        abortControllers.delete(addedAbortKey);
+      });
+      return;
+    }
     abortControllers.set(abortKey, { abortController, ...endpointOption });
 
     res.on('finish', function () {

api/server/routes/ask/gptPlugins.js

@@ -148,15 +148,6 @@ router.post(
       }
     };
 
-    const onChainEnd = () => {
-      saveMessage({ ...userMessage, user });
-      sendIntermediateMessage(res, {
-        plugins,
-        parentMessageId: userMessage.messageId,
-        messageId: responseMessageId,
-      });
-    };
-
     const getAbortData = () => ({
       sender,
       conversationId,
@@ -167,12 +158,23 @@ router.post(
       userMessage,
       promptTokens,
     });
-    const { abortController, onStart } = createAbortController(req, res, getAbortData);
+    const { abortController, onStart } = createAbortController(req, res, getAbortData, getReqData);
 
     try {
       endpointOption.tools = await validateTools(user, endpointOption.tools);
       const { client } = await initializeClient({ req, res, endpointOption });
 
+      const onChainEnd = () => {
+        if (!client.skipSaveUserMessage) {
+          saveMessage({ ...userMessage, user });
+        }
+        sendIntermediateMessage(res, {
+          plugins,
+          parentMessageId: userMessage.messageId,
+          messageId: responseMessageId,
+        });
+      };
+
       let response = await client.sendMessage(text, {
         user,
         conversationId,

api/server/routes/edit/gptPlugins.js

@@ -103,21 +103,6 @@ router.post(
       },
     });
 
-    const onAgentAction = (action, start = false) => {
-      const formattedAction = formatAction(action);
-      plugin.inputs.push(formattedAction);
-      plugin.latest = formattedAction.plugin;
-      if (!start) {
-        saveMessage({ ...userMessage, user });
-      }
-      sendIntermediateMessage(res, {
-        plugin,
-        parentMessageId: userMessage.messageId,
-        messageId: responseMessageId,
-      });
-      // logger.debug('PLUGIN ACTION', formattedAction);
-    };
-
     const onChainEnd = (data) => {
       let { intermediateSteps: steps } = data;
       plugin.outputs = steps && steps[0].action ? formatSteps(steps) : 'An error occurred.';
@@ -141,12 +126,27 @@ router.post(
       userMessage,
       promptTokens,
     });
-    const { abortController, onStart } = createAbortController(req, res, getAbortData);
+    const { abortController, onStart } = createAbortController(req, res, getAbortData, getReqData);
 
     try {
       endpointOption.tools = await validateTools(user, endpointOption.tools);
       const { client } = await initializeClient({ req, res, endpointOption });
 
+      const onAgentAction = (action, start = false) => {
+        const formattedAction = formatAction(action);
+        plugin.inputs.push(formattedAction);
+        plugin.latest = formattedAction.plugin;
+        if (!start && !client.skipSaveUserMessage) {
+          saveMessage({ ...userMessage, user });
+        }
+        sendIntermediateMessage(res, {
+          plugin,
+          parentMessageId: userMessage.messageId,
+          messageId: responseMessageId,
+        });
+        // logger.debug('PLUGIN ACTION', formattedAction);
+      };
+
       let response = await client.sendMessage(text, {
         user,
         generation,