🚀 feat: Enhance Message Editing with File Resubmission (#2347)

* chore: fix type issue with File Table fakeData

* refactor: new lazy loading image strategy and load images/files as part of Message Container

* feat: resubmit files when editing messages with attached files

client/src/common/types.ts

@@ -181,6 +181,7 @@ export type TAskProps = {
 export type TOptions = {
   editedMessageId?: string | null;
   editedText?: string | null;
+  resubmitFiles?: boolean;
   isRegenerate?: boolean;
   isContinued?: boolean;
   isEdited?: boolean;

client/src/components/Chat/Input/Files/Table/fakeData.ts

@@ -17,6 +17,7 @@ export const files: TFile[] = [
     updatedAt: '2024-01-23T18:25:48.153Z',
     usage: 0,
     user: '652ac880c4102a77fe54c5db',
+    embedded: false,
   },
   {
     _id: '65b004abd70ce86b9146e861',
@@ -34,6 +35,7 @@ export const files: TFile[] = [
     usage: 0,
     user: '652ac880c4102a77fe54c5db',
     width: 1024,
+    embedded: false,
   },
   {
     _id: '65b00495d70ce86b9146adc1',
@@ -51,6 +53,7 @@ export const files: TFile[] = [
     usage: 0,
     user: '652ac880c4102a77fe54c5db',
     width: 1024,
+    embedded: false,
   },
   {
     _id: '65b00494d70ce86b9146ace6',
@@ -68,5 +71,6 @@ export const files: TFile[] = [
     usage: 0,
     user: '652ac880c4102a77fe54c5db',
     width: 1024,
+    embedded: false,
   },
 ];

client/src/components/Chat/Messages/Content/Container.tsx

@@ -1,6 +1,9 @@
-// Container Component
-const Container = ({ children }: { children: React.ReactNode }) => (
+import { TMessage } from 'librechat-data-provider';
+import Files from './Files';
+
+const Container = ({ children, message }: { children: React.ReactNode; message: TMessage }) => (
   <div className="text-message flex min-h-[20px] flex-col items-start gap-3 overflow-x-auto [.text-message+&]:mt-5">
+    {message.isCreatedByUser && <Files message={message} />}
     {children}
   </div>
 );

client/src/components/Chat/Messages/Content/ContentParts.tsx

@@ -38,7 +38,7 @@ any) => {
         {!isSubmitting && unfinished && (
           <Suspense>
             <DelayedRender delay={250}>
-              <UnfinishedMessage key={`unfinished-${messageId}`} />
+              <UnfinishedMessage message={message} key={`unfinished-${messageId}`} />
             </DelayedRender>
           </Suspense>
         )}

client/src/components/Chat/Messages/Content/EditMessage.tsx

@@ -3,10 +3,10 @@ import { EModelEndpoint } from 'librechat-data-provider';
 import { useState, useRef, useEffect, useCallback } from 'react';
 import { useUpdateMessageMutation } from 'librechat-data-provider/react-query';
 import type { TEditProps } from '~/common';
-import Container from '~/components/Messages/Content/Container';
 import { cn, removeFocusOutlines } from '~/utils';
 import { useChatContext } from '~/Providers';
 import { useLocalize } from '~/hooks';
+import Container from './Container';
 
 const EditMessage = ({
   text,
@@ -39,11 +39,16 @@ const EditMessage = ({
 
   const resubmitMessage = () => {
     if (message.isCreatedByUser) {
-      ask({
-        text: editedText,
-        parentMessageId,
-        conversationId,
-      });
+      ask(
+        {
+          text: editedText,
+          parentMessageId,
+          conversationId,
+        },
+        {
+          resubmitFiles: true,
+        },
+      );
 
       setSiblingIdx((siblingIdx ?? 0) - 1);
     } else {
@@ -105,7 +110,7 @@ const EditMessage = ({
   );
 
   return (
-    <Container>
+    <Container message={message}>
       <TextareaAutosize
         ref={textAreaRef}
         onChange={(e) => {

client/src/components/Chat/Messages/Content/Files.tsx

@@ -0,0 +1,39 @@
+import { useMemo, memo } from 'react';
+import type { TFile, TMessage } from 'librechat-data-provider';
+import FileContainer from '~/components/Chat/Input/Files/FileContainer';
+import Image from './Image';
+
+const Files = ({ message }: { message: TMessage }) => {
+  const imageFiles = useMemo(() => {
+    return message?.files?.filter((file) => file.type?.startsWith('image/')) || [];
+  }, [message?.files]);
+
+  const otherFiles = useMemo(() => {
+    return message?.files?.filter((file) => !file.type?.startsWith('image/')) || [];
+  }, [message?.files]);
+
+  return (
+    <>
+      {otherFiles.length > 0 &&
+        otherFiles.map((file) => <FileContainer key={file.file_id} file={file as TFile} />)}
+      {imageFiles &&
+        imageFiles.map((file) => (
+          <Image
+            key={file.file_id}
+            imagePath={file?.preview ?? file.filepath ?? ''}
+            height={file.height ?? 1920}
+            width={file.width ?? 1080}
+            altText={file.filename ?? 'Uploaded Image'}
+            placeholderDimensions={{
+              height: `${file.height ?? 1920}px`,
+              width: `${file.height ?? 1080}px`,
+            }}
+            // n={imageFiles.length}
+            // i={i}
+          />
+        ))}
+    </>
+  );
+};
+
+export default memo(Files);

client/src/components/Chat/Messages/Content/Image.tsx

@@ -1,18 +1,35 @@
-import React, { useState, useEffect } from 'react';
+import React, { useState, useRef, useMemo } from 'react';
 import { LazyLoadImage } from 'react-lazy-load-image-component';
 import * as Dialog from '@radix-ui/react-dialog';
 import DialogImage from './DialogImage';
 import { cn } from '~/utils';
 
+const scaleImage = ({
+  originalWidth,
+  originalHeight,
+  containerRef,
+}: {
+  originalWidth: number;
+  originalHeight: number;
+  containerRef: React.RefObject<HTMLDivElement>;
+}) => {
+  const containerWidth = containerRef.current?.offsetWidth ?? 0;
+  if (containerWidth === 0 || originalWidth === undefined || originalHeight === undefined) {
+    return { width: 'auto', height: 'auto' };
+  }
+  const aspectRatio = originalWidth / originalHeight;
+  const scaledWidth = Math.min(containerWidth, originalWidth);
+  const scaledHeight = scaledWidth / aspectRatio;
+  return { width: `${scaledWidth}px`, height: `${scaledHeight}px` };
+};
+
 const Image = ({
   imagePath,
   altText,
   height,
   width,
   placeholderDimensions,
-}: // n,
-// i,
-{
+}: {
   imagePath: string;
   altText: string;
   height: number;
@@ -21,69 +38,49 @@ const Image = ({
     height: string;
     width: string;
   };
-  // n: number;
-  // i: number;
 }) => {
   const [isLoaded, setIsLoaded] = useState(false);
+  const containerRef = useRef<HTMLDivElement>(null);
+
   const handleImageLoad = () => setIsLoaded(true);
-  const [minDisplayTimeElapsed, setMinDisplayTimeElapsed] = useState(false);
 
-  useEffect(() => {
-    let timer: NodeJS.Timeout;
-    if (isLoaded) {
-      timer = setTimeout(() => setMinDisplayTimeElapsed(true), 150);
-    }
-    return () => clearTimeout(timer);
-  }, [isLoaded]);
-  // const makeSquare = n >= 3 && i < 2;
-
-  let placeholderHeight = '288px';
-  if (placeholderDimensions?.height && placeholderDimensions?.width) {
-    placeholderHeight = placeholderDimensions.height;
-  } else if (height > width) {
-    placeholderHeight = '900px';
-  } else if (height === width) {
-    placeholderHeight = width + 'px';
-  }
+  const { width: scaledWidth, height: scaledHeight } = useMemo(
+    () =>
+      scaleImage({
+        originalWidth: Number(placeholderDimensions?.width?.split('px')[0]) ?? width,
+        originalHeight: Number(placeholderDimensions?.height?.split('px')[0]) ?? height,
+        containerRef,
+      }),
+    [placeholderDimensions, height, width],
+  );
 
   return (
     <Dialog.Root>
-      <div className="">
+      <div ref={containerRef}>
         <div className="relative mt-1 flex h-auto w-full max-w-lg items-center justify-center overflow-hidden bg-gray-200 text-gray-500 dark:bg-gray-700 dark:text-gray-400">
           <Dialog.Trigger asChild>
             <button type="button" aria-haspopup="dialog" aria-expanded="false">
               <LazyLoadImage
-                // loading="lazy"
                 alt={altText}
                 onLoad={handleImageLoad}
                 visibleByDefault={true}
                 className={cn(
-                  'max-h-[900px] max-w-full opacity-100 transition-opacity duration-300',
-                  // n >= 3 && i < 2 ? 'aspect-square object-cover' : '',
-                  isLoaded && minDisplayTimeElapsed ? 'opacity-100' : 'opacity-0',
+                  'opacity-100 transition-opacity duration-100',
+                  isLoaded ? 'opacity-100' : 'opacity-0',
                 )}
                 src={imagePath}
                 style={{
-                  height: isLoaded && minDisplayTimeElapsed ? 'auto' : placeholderHeight,
-                  width: placeholderDimensions?.width ?? width,
+                  width: scaledWidth,
+                  height: 'auto',
                   color: 'transparent',
                 }}
-                placeholder={
-                  <div
-                    style={{
-                      height: isLoaded && minDisplayTimeElapsed ? 'auto' : placeholderHeight,
-                      width: placeholderDimensions?.width ?? width,
-                    }}
-                  />
-                }
+                placeholder={<div style={{ width: scaledWidth, height: scaledHeight }} />}
               />
             </button>
           </Dialog.Trigger>
         </div>
       </div>
-      {isLoaded && minDisplayTimeElapsed && (
-        <DialogImage src={imagePath} height={height} width={width} />
-      )}
+      {isLoaded && <DialogImage src={imagePath} height={height} width={width} />}
     </Dialog.Root>
   );
 };

client/src/components/Chat/Messages/Content/MessageContent.tsx

@@ -1,7 +1,6 @@
 import { Fragment, Suspense } from 'react';
-import type { TResPlugin, TFile } from 'librechat-data-provider';
-import type { TMessageContentProps, TText, TDisplayProps } from '~/common';
-import FileContainer from '~/components/Chat/Input/Files/FileContainer';
+import type { TMessage, TResPlugin } from 'librechat-data-provider';
+import type { TMessageContentProps, TDisplayProps } from '~/common';
 import Plugin from '~/components/Messages/Content/Plugin';
 import Error from '~/components/Messages/Content/Error';
 import { DelayedRender } from '~/components/ui';
@@ -9,11 +8,14 @@ import EditMessage from './EditMessage';
 import Container from './Container';
 import Markdown from './Markdown';
 import { cn } from '~/utils';
-import Image from './Image';
 
-export const ErrorMessage = ({ text, className = '' }: TText) => {
+export const ErrorMessage = ({
+  text,
+  message,
+  className = '',
+}: Pick<TDisplayProps, 'text' | 'className' | 'message'>) => {
   return (
-    <Container>
+    <Container message={message}>
       <div
         className={cn(
           'rounded-md border border-red-500 bg-red-500/10 px-3 py-2 text-sm text-gray-600 dark:text-gray-200',
@@ -28,31 +30,8 @@ export const ErrorMessage = ({ text, className = '' }: TText) => {
 
 // Display Message Component
 const DisplayMessage = ({ text, isCreatedByUser, message, showCursor }: TDisplayProps) => {
-  const files: TFile[] = [];
-  const imageFiles = message?.files
-    ? message.files.filter((file) => {
-      if (file.type && file.type.startsWith('image/')) {
-        return true;
-      }
-
-      files.push(file);
-    })
-    : null;
   return (
-    <Container>
-      {files.length > 0 && files.map((file) => <FileContainer key={file.file_id} file={file} />)}
-      {imageFiles &&
-        imageFiles.map((file) => (
-          <Image
-            key={file.file_id}
-            imagePath={file?.preview ?? file.filepath ?? ''}
-            height={file.height ?? 1920}
-            width={file.width ?? 1080}
-            altText={file.filename ?? 'Uploaded Image'}
-            // n={imageFiles.length}
-            // i={i}
-          />
-        ))}
+    <Container message={message}>
       <div
         className={cn(
           showCursor && !!text?.length ? 'result-streaming' : '',
@@ -71,8 +50,11 @@ const DisplayMessage = ({ text, isCreatedByUser, message, showCursor }: TDisplay
 };
 
 // Unfinished Message Component
-export const UnfinishedMessage = () => (
-  <ErrorMessage text="The response is incomplete; it's either still processing, was cancelled, or censored. Refresh or try a different prompt." />
+export const UnfinishedMessage = ({ message }: { message: TMessage }) => (
+  <ErrorMessage
+    message={message}
+    text="The response is incomplete; it's either still processing, was cancelled, or censored. Refresh or try a different prompt."
+  />
 );
 
 // Content Component
@@ -86,7 +68,7 @@ const MessageContent = ({
   ...props
 }: TMessageContentProps) => {
   if (error) {
-    return <ErrorMessage text={text} />;
+    return <ErrorMessage message={props.message} text={text} />;
   } else if (edit) {
     return <EditMessage text={text} isSubmitting={isSubmitting} {...props} />;
   } else {
@@ -143,7 +125,7 @@ const MessageContent = ({
           {!isSubmitting && unfinished && (
             <Suspense>
               <DelayedRender delay={250}>
-                <UnfinishedMessage key={`unfinished-${messageId}-${idx}`} />
+                <UnfinishedMessage message={message} key={`unfinished-${messageId}-${idx}`} />
               </DelayedRender>
             </Suspense>
           )}

client/src/components/Chat/Messages/Content/Part.tsx

@@ -53,11 +53,11 @@ export default function Part({
   }
 
   if (part.type === ContentTypes.ERROR) {
-    return <ErrorMessage text={part[ContentTypes.TEXT].value} className="my-2" />;
+    return <ErrorMessage message={message} text={part[ContentTypes.TEXT].value} className="my-2" />;
   } else if (part.type === ContentTypes.TEXT) {
     // Access the value property
     return (
-      <Container>
+      <Container message={message}>
         <div className="markdown prose dark:prose-invert light dark:text-gray-70 my-1 w-full break-words">
           <DisplayMessage
             text={part[ContentTypes.TEXT].value}
@@ -104,7 +104,7 @@ export default function Part({
     if (isImageVisionTool(toolCall)) {
       if (isSubmitting && showCursor) {
         return (
-          <Container>
+          <Container message={message}>
             <div className="markdown prose dark:prose-invert light dark:text-gray-70 my-1 w-full break-words">
               <DisplayMessage
                 text={''}

client/src/hooks/useChatHelpers.ts

@@ -94,6 +94,7 @@ export default function useChatHelpers(index = 0, paramId: string | undefined) {
     {
       editedText = null,
       editedMessageId = null,
+      resubmitFiles = false,
       isRegenerate = false,
       isContinued = false,
       isEdited = false,
@@ -177,7 +178,7 @@ export default function useChatHelpers(index = 0, paramId: string | undefined) {
       error: false,
     };
 
-    const reuseFiles = isRegenerate && parentMessage?.files;
+    const reuseFiles = (isRegenerate || resubmitFiles) && parentMessage?.files;
     if (reuseFiles && parentMessage.files?.length) {
       currentMsg.files = parentMessage.files;
       setFiles(new Map());