feat: Edit AI Messages, Edit Messages in Place (#825)

* refactor: replace lodash import with specific function import

fix(api): esm imports to cjs

* refactor(Messages.tsx): convert to TS, out-source scrollToDiv logic to a custom hook
fix(ScreenshotContext.tsx): change Ref to RefObject in ScreenshotContextType
feat(useScrollToRef.ts): add useScrollToRef hook for scrolling to a ref with throttle
fix(Chat.tsx): update import path for Messages component
fix(Search.tsx): update import path for Messages component

* chore(types.ts): add TAskProps and TOptions types
refactor(useMessageHandler.ts): use TAskFunction type for ask function signature

* refactor(Message/Content): convert to TS, move Plugin component to Content dir

* feat(MessageContent.tsx): add MessageContent component for displaying and editing message content
feat(index.ts): export MessageContent component from Messages/Content directory

* wip(Message.jsx): conversion and use of new component in progress

* refactor: convert Message.jsx to TS and fix typing/imports based on changes

* refactor: add typed props and refactor MultiMessage to TS, fix typing issues resulting from the conversion

* edit message in progress

* feat: complete edit AI message logic, refactor continue logic

* feat(middleware): add validateMessageReq middleware
feat(routes): add validation for message requests using validateMessageReq middleware
feat(routes): add create, read, update, and delete routes for messages

* feat: complete frontend logic for editing messages in place
feat(messages.js): update route for updating a specific message
- Change the route for updating a message to include the messageId in the URL
- Update the request handler to use the messageId from the request parameters and the text from the request body
- Call the updateMessage function with the updated parameters

feat(MessageContent.tsx): add functionality to update a message
- Import the useUpdateMessageMutation hook from the data provider
- Destructure the conversationId, parentMessageId, and messageId from the message object
- Create a mutation function using the useUpdateMessageMutation hook
- Implement the updateMessage function to call the mutation function with the updated message parameters
- Update the messages state to reflect the updated message text

feat(api-endpoints.ts): update messages endpoint to include messageId
- Update the messages endpoint to include the messageId as an optional parameter

feat(data-service.ts): add updateMessage function
- Implement the updateMessage function to make a PUT request to

* fix(messages.js): make updateMessage function asynchronous and await its execution

* style(EditIcon): make icon active for AI message

* feat(gptPlugins/anthropic): add edit support

* fix(validateMessageReq.js): handle case when conversationId is 'new' and return empty array
feat(Message.tsx): pass message prop to SiblingSwitch component
refactor(SiblingSwitch.tsx): convert to TS

* fix(useMessageHandler.ts): remove message from currentMessages if isContinued is true
feat(useMessageHandler.ts): add support for submission messages in setMessages
fix(useServerStream.ts): remove unnecessary conditional in setMessages
fix(useServerStream.ts): remove isContinued variable from submission

* fix(continue): switch to continued message generation when continuing an earlier branch in conversation

* fix(abortMiddleware.js): fix condition to check partialText length
chore(abortMiddleware.js): add error logging when abortMessage fails

* refactor(MessageHeader.tsx): convert to TS
fix(Plugin.tsx): add default value for className prop in Plugin component

* refactor(MultiMessage.tsx): remove commented out code
docs(MultiMessage.tsx): update comment to clarify when siblingIdx is reset

* fix(GenerationButtons): optimistic state for continue button

* fix(MessageContent.tsx): add data-testid attribute to message text editor
fix(messages.spec.ts): update waitForServerStream function to include edit endpoint check
feat(messages.spec.ts): add test case for editing messages

* fix(HoverButtons & Message & useGenerations): Refactor edit functionality and related conditions

- Update enterEdit function signature and prop
- Create and utilize hideEditButton variable
- Enhance conditions for edit button visibility and active state
- Update button event handlers
- Introduce isEditableEndpoint in useGenerations and refine continueSupported condition.

* fix(useGenerations.ts): fix condition for hideEditButton to include error and searchResult
chore(data-provider): bump version to 0.1.6
fix(types.ts): add status property to TError type

* chore: bump @dqbd/tiktoken to 1.0.7

* fix(abortMiddleware.js): add required isCreatedByUser property to the error response object

* refactor(Message.tsx): remove unnecessary props from SiblingSwitch component, as setLatestMessage is firing on every switch already
refactor(SiblingSwitch.tsx): remove unused imports and code

* chore(BaseClient.js): move console.debug statements back inside if block

api/server/middleware/abortMiddleware.js

@@ -79,6 +79,7 @@ const handleAbortError = async (res, req, error, data) => {
       cancelled: false,
       error: true,
       text: error.message,
+      isCreatedByUser: false,
     };
     if (abortControllers.has(conversationId)) {
       const { abortController } = abortControllers.get(conversationId);
@@ -89,10 +90,11 @@ const handleAbortError = async (res, req, error, data) => {
     handleError(res, errorMessage);
   };
 
-  if (partialText?.length > 2) {
+  if (partialText && partialText.length > 5) {
     try {
       return await abortMessage(req, res);
     } catch (err) {
+      console.error(err);
       return respondWithError();
     }
   } else {

api/server/middleware/index.js

@@ -3,6 +3,7 @@ const setHeaders = require('./setHeaders');
 const requireJwtAuth = require('./requireJwtAuth');
 const requireLocalAuth = require('./requireLocalAuth');
 const validateEndpoint = require('./validateEndpoint');
+const validateMessageReq = require('./validateMessageReq');
 const buildEndpointOption = require('./buildEndpointOption');
 const validateRegistration = require('./validateRegistration');
 
@@ -12,6 +13,7 @@ module.exports = {
   requireJwtAuth,
   requireLocalAuth,
   validateEndpoint,
+  validateMessageReq,
   buildEndpointOption,
   validateRegistration,
 };

api/server/middleware/validateMessageReq.js

@@ -0,0 +1,28 @@
+const { getConvo } = require('../../models');
+
+// Middleware to validate conversationId and user relationship
+const validateMessageReq = async (req, res, next) => {
+  let conversationId = req.params.conversationId || req.body.conversationId;
+
+  if (conversationId === 'new') {
+    return res.status(200).send([]);
+  }
+
+  if (!conversationId && req.body.message) {
+    conversationId = req.body.message.conversationId;
+  }
+
+  const conversation = await getConvo(req.user.id, conversationId);
+
+  if (!conversation) {
+    return res.status(404).json({ error: 'Conversation not found' });
+  }
+
+  if (conversation.user !== req.user.id) {
+    return res.status(403).json({ error: 'User not authorized for this conversation' });
+  }
+
+  next();
+};
+
+module.exports = validateMessageReq;

api/server/routes/edit/anthropic.js

@@ -29,11 +29,12 @@ router.post(
       endpointOption,
       conversationId,
       responseMessageId,
+      isContinued = false,
       parentMessageId = null,
       overrideParentMessageId = null,
     } = req.body;
     console.log('edit log');
-    console.dir({ text, conversationId, endpointOption }, { depth: null });
+    console.dir({ text, generation, isContinued, conversationId, endpointOption }, { depth: null });
     let metadata;
     let userMessage;
     let lastSavedTimestamp = 0;
@@ -41,7 +42,10 @@ router.post(
     const userMessageId = parentMessageId;
 
     const addMetadata = (data) => (metadata = data);
-    const getIds = (data) => (userMessage = data.userMessage);
+    const getIds = (data) => {
+      userMessage = data.userMessage;
+      responseMessageId = data.responseMessageId;
+    };
 
     const { onProgress: progressCallback, getPartialText } = createOnProgress({
       generation,
@@ -87,6 +91,8 @@ router.post(
 
       let response = await client.sendMessage(text, {
         user: req.user.id,
+        generation,
+        isContinued,
         isEdited: true,
         conversationId,
         parentMessageId,

api/server/routes/edit/gptPlugins.js

@@ -30,11 +30,12 @@ router.post(
       endpointOption,
       conversationId,
       responseMessageId,
+      isContinued = false,
       parentMessageId = null,
       overrideParentMessageId = null,
     } = req.body;
     console.log('edit log');
-    console.dir({ text, conversationId, endpointOption }, { depth: null });
+    console.dir({ text, generation, isContinued, conversationId, endpointOption }, { depth: null });
     let metadata;
     let userMessage;
     let lastSavedTimestamp = 0;
@@ -50,7 +51,10 @@ router.post(
     };
 
     const addMetadata = (data) => (metadata = data);
-    const getIds = (data) => (userMessage = data.userMessage);
+    const getIds = (data) => {
+      userMessage = data.userMessage;
+      responseMessageId = data.responseMessageId;
+    };
 
     const {
       onProgress: progressCallback,
@@ -128,6 +132,8 @@ router.post(
 
       let response = await client.sendMessage(text, {
         user,
+        generation,
+        isContinued,
         isEdited: true,
         conversationId,
         parentMessageId,

api/server/routes/edit/openAI.js

@@ -29,11 +29,12 @@ router.post(
       endpointOption,
       conversationId,
       responseMessageId,
+      isContinued = false,
       parentMessageId = null,
       overrideParentMessageId = null,
     } = req.body;
     console.log('edit log');
-    console.dir({ text, conversationId, endpointOption }, { depth: null });
+    console.dir({ text, generation, isContinued, conversationId, endpointOption }, { depth: null });
     let metadata;
     let userMessage;
     let lastSavedTimestamp = 0;
@@ -41,7 +42,10 @@ router.post(
     const userMessageId = parentMessageId;
 
     const addMetadata = (data) => (metadata = data);
-    const getIds = (data) => (userMessage = data.userMessage);
+    const getIds = (data) => {
+      userMessage = data.userMessage;
+      responseMessageId = data.responseMessageId;
+    };
 
     const { onProgress: progressCallback, getPartialText } = createOnProgress({
       generation,
@@ -90,6 +94,8 @@ router.post(
 
       let response = await client.sendMessage(text, {
         user: req.user.id,
+        generation,
+        isContinued,
         isEdited: true,
         conversationId,
         parentMessageId,

api/server/routes/messages.js

@@ -1,11 +1,50 @@
 const express = require('express');
 const router = express.Router();
-const { getMessages } = require('../../models/Message');
-const requireJwtAuth = require('../middleware/requireJwtAuth');
+const {
+  getMessages,
+  updateMessage,
+  saveConvo,
+  saveMessage,
+  deleteMessages,
+} = require('../../models');
+const { requireJwtAuth, validateMessageReq } = require('../middleware/');
 
-router.get('/:conversationId', requireJwtAuth, async (req, res) => {
+router.get('/:conversationId', requireJwtAuth, validateMessageReq, async (req, res) => {
   const { conversationId } = req.params;
   res.status(200).send(await getMessages({ conversationId }));
 });
 
+// CREATE
+router.post('/:conversationId', requireJwtAuth, validateMessageReq, async (req, res) => {
+  const message = req.body;
+  const savedMessage = await saveMessage(message);
+  await saveConvo(req.user.id, savedMessage);
+  res.status(201).send(savedMessage);
+});
+
+// READ
+router.get('/:conversationId/:messageId', requireJwtAuth, validateMessageReq, async (req, res) => {
+  const { conversationId, messageId } = req.params;
+  res.status(200).send(await getMessages({ conversationId, messageId }));
+});
+
+// UPDATE
+router.put('/:conversationId/:messageId', requireJwtAuth, validateMessageReq, async (req, res) => {
+  const { messageId } = req.params;
+  const { text } = req.body;
+  res.status(201).send(await updateMessage({ messageId, text }));
+});
+
+// DELETE
+router.delete(
+  '/:conversationId/:messageId',
+  requireJwtAuth,
+  validateMessageReq,
+  async (req, res) => {
+    const { messageId } = req.params;
+    await deleteMessages({ messageId });
+    res.status(204).send();
+  },
+);
+
 module.exports = router;

api/server/utils/handleText.js

@@ -1,4 +1,4 @@
-const _ = require('lodash');
+const partialRight = require('lodash/partialRight');
 const citationRegex = /\[\^\d+?\^]/g;
 const { getCitations, citeText } = require('./citations');
 const cursor = '<span className="result-streaming">█</span>';
@@ -73,7 +73,7 @@ const createOnProgress = ({ generation = '', onProgress: _onProgress }) => {
   };
 
   const onProgress = (opts) => {
-    return _.partialRight(progressCallback, opts);
+    return partialRight(progressCallback, opts);
   };
 
   const getPartialText = () => {