fix: pass running to message context (#1859)

* fix: pass running to message context * fix: attachment
Chainlit · Feb 5, 2025 · c981730 · c981730
1 parent 44d94bd
commit c981730
Show file tree

Hide file tree

Showing 10 changed files with 92 additions and 39 deletions.
diff --git a/backend/chainlit/element.py b/backend/chainlit/element.py
@@ -166,13 +166,9 @@ def from_dict(cls, e_dict: ElementDict):
             "mime": mime_type,
         }
 
-        # Image handling (excluding SVG which is treated as a file)
-        if type == "image" and "svg" not in mime_type:
+        if type == "image":
             return Image(size="medium", **common_params)  # type: ignore[arg-type]
 
-        elif type == "pdf":
-            return Pdf(page=e_dict.get("page"), **common_params)  # type: ignore[arg-type]
-
         elif type == "audio":
             return Audio(auto_play=e_dict.get("autoPlay", False), **common_params)  # type: ignore[arg-type]
 
@@ -182,15 +178,9 @@ def from_dict(cls, e_dict: ElementDict):
                 **common_params,  # type: ignore[arg-type]
             )
 
-        elif type == "text":
-            return Text(language=e_dict.get("language"), **common_params)  # type: ignore[arg-type]
-
-        elif type == "tasklist":
-            return TaskList(**common_params)  # type: ignore[arg-type]
         elif type == "plotly":
             return Plotly(size=e_dict.get("size", "medium"), **common_params)  # type: ignore[arg-type]
-        elif type == "dataframe":
-            return Dataframe(size=e_dict.get("size", "large"), **common_params)  # type: ignore[arg-type]
+
         elif type == "custom":
             return CustomElement(props=e_dict.get("props", {}), **common_params)  # type: ignore[arg-type]
         else:

diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
@@ -41,15 +41,20 @@ function App() {
     if (!isAuthenticated || !isReady || !chatProfileOk) {
       return;
     }
-    
+
     connect({
       transports: window.transports,
       userEnv
     });
   }, [userEnv, isAuthenticated, connect, isReady, chatProfileOk]);
 
   useEffect(() => {
-    if (!configLoaded || !config || !config.chatProfiles?.length || chatProfile) {
+    if (
+      !configLoaded ||
+      !config ||
+      !config.chatProfiles?.length ||
+      chatProfile
+    ) {
       return;
     }
 

diff --git a/frontend/src/components/Elements/Image.tsx b/frontend/src/components/Elements/Image.tsx
@@ -22,7 +22,7 @@ const ImageElement = ({ element }: { element: IImageElement }) => {
 
   return (
     <>
-      <div className="rounded-sm bg-muted overflow-hidden">
+      <div className="rounded-sm bg-accent overflow-hidden">
         <img
           className={cn(
             'mx-auto block max-w-full h-auto',

diff --git a/frontend/src/components/QuiltedGrid.tsx b/frontend/src/components/QuiltedGrid.tsx
@@ -26,11 +26,34 @@ const QuiltedGrid = <T extends IImageElement | IVideoElement>({
   renderElement: Renderer,
   className
 }: QuiltedGridProps<T>) => {
+  // If there's only one element, use a simpler layout
+  if (elements.length === 1) {
+    const element = elements[0];
+    const size = sizeToUnit(element);
+
+    return (
+      <div
+        className={cn(
+          'w-full',
+          // Adjust max-width based on size
+          size === 1
+            ? 'max-w-[150px]'
+            : size === 2
+            ? 'max-w-[300px]'
+            : 'max-w-[600px]',
+          className
+        )}
+      >
+        <Renderer element={element} />
+      </div>
+    );
+  }
+
   return (
     <div
       className={cn(
         'grid grid-cols-4 gap-2 w-full max-w-[600px]',
-        'transform-gpu', // equivalent to transform: translateZ(0)
+        'transform-gpu',
         className
       )}
     >

diff --git a/frontend/src/components/chat/MessageComposer/Attachment.tsx b/frontend/src/components/chat/MessageComposer/Attachment.tsx
@@ -2,6 +2,12 @@ import React from 'react';
 import { DefaultExtensionType, FileIcon, defaultStyles } from 'react-file-icon';
 
 import { Card } from '@/components/ui/card';
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipProvider,
+  TooltipTrigger
+} from '@/components/ui/tooltip';
 
 interface AttachmentProps {
   name: string;
@@ -15,17 +21,26 @@ const Attachment: React.FC<AttachmentProps> = ({ name, mime, children }) => {
   ) as DefaultExtensionType;
 
   return (
-    <div className="relative h-[58px]">
-      {children}
-      <Card className="h-full p-2 flex flex-row items-center gap-3 rounded-lg w-full max-w-[200px] border">
-        <div className="w-10">
-          <FileIcon {...defaultStyles[extension]} extension={extension} />
-        </div>
-        <span className="truncate w-[80%] font-medium text-sm font-medium">
-          {name}
-        </span>
-      </Card>
-    </div>
+    <TooltipProvider delayDuration={100}>
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <div className="relative h-[58px]">
+            {children}
+            <Card className="h-full p-2 flex flex-row items-center gap-3 rounded-lg w-full max-w-[200px] border">
+              <div className="w-10">
+                <FileIcon {...defaultStyles[extension]} extension={extension} />
+              </div>
+              <span className="truncate w-[80%] font-medium text-sm font-medium">
+                {name}
+              </span>
+            </Card>
+          </div>
+        </TooltipTrigger>
+        <TooltipContent>
+          <p>{name}</p>
+        </TooltipContent>
+      </Tooltip>
+    </TooltipProvider>
   );
 };
 

diff --git a/frontend/src/components/chat/Messages/Message/Content/InlinedElements/InlinedPDFList.tsx b/frontend/src/components/chat/Messages/Message/Content/InlinedElements/InlinedPDFList.tsx
@@ -10,7 +10,7 @@ const InlinedPDFList = ({ items }: Props) => (
   <div className="flex flex-col gap-2">
     {items.map((pdf, i) => {
       return (
-        <div key={i} className="w-[90%] h-[400px]">
+        <div key={i} className="h-[400px]">
           <PDFElement element={pdf} />
         </div>
       );

diff --git a/frontend/src/components/chat/Messages/Message/Content/InlinedElements/index.tsx b/frontend/src/components/chat/Messages/Message/Content/InlinedElements/index.tsx
@@ -1,3 +1,5 @@
+import { cn } from '@/lib/utils';
+
 import type { ElementType, IMessageElement } from '@chainlit/react-client';
 
 import { InlinedCustomElementList } from './InlineCustomElementList';
@@ -12,9 +14,10 @@ import { InlinedVideoList } from './InlinedVideoList';
 
 interface Props {
   elements: IMessageElement[];
+  className?: string;
 }
 
-const InlinedElements = ({ elements }: Props) => {
+const InlinedElements = ({ elements, className }: Props) => {
   if (!elements.length) {
     return null;
   }
@@ -42,7 +45,7 @@ const InlinedElements = ({ elements }: Props) => {
   );
 
   return (
-    <div className="flex flex-col gap-4">
+    <div className={cn('flex flex-col gap-4', className)}>
       {elementsByType.custom?.length ? (
         <InlinedCustomElementList items={elementsByType.custom} />
       ) : null}

diff --git a/frontend/src/components/chat/Messages/Message/UserMessage.tsx b/frontend/src/components/chat/Messages/Message/UserMessage.tsx
@@ -1,8 +1,10 @@
+import { cn } from '@/lib/utils';
 import { MessageContext } from 'contexts/MessageContext';
-import { useContext, useState } from 'react';
+import { useContext, useMemo, useState } from 'react';
 import { useSetRecoilState } from 'recoil';
 
 import {
+  IMessageElement,
   IStep,
   messagesState,
   useChatInteract,
@@ -14,12 +16,16 @@ import { Pencil } from '@/components/icons/Pencil';
 import { Button } from '@/components/ui/button';
 import { Translator } from 'components/i18n';
 
+import { InlinedElements } from './Content/InlinedElements';
+
 interface Props {
   message: IStep;
+  elements: IMessageElement[];
 }
 
 export default function UserMessage({
   message,
+  elements,
   children
 }: React.PropsWithChildren<Props>) {
   const config = useConfig();
@@ -30,6 +36,12 @@ export default function UserMessage({
   const [isEditing, setIsEditing] = useState(false);
   const [editValue, setEditValue] = useState('');
 
+  const inlineElements = useMemo(() => {
+    return elements.filter(
+      (el) => el.forId === message.id && el.display === 'inline'
+    );
+  }, [message.id, elements]);
+
   const isEditable = !!config.config?.features.edit_message;
 
   const handleEdit = () => {
@@ -49,7 +61,9 @@ export default function UserMessage({
   };
 
   return (
-    <div className="flex flex-col w-full">
+    <div className="flex flex-col w-full gap-1">
+      <InlinedElements elements={inlineElements} className="items-end" />
+
       <div className="flex flex-row items-center gap-1 w-full group">
         {!isEditing && isEditable && (
           <Button
@@ -66,13 +80,15 @@ export default function UserMessage({
           </Button>
         )}
         <div
-          className={`px-5 py-2.5 relative bg-accent rounded-3xl
-            ${isEditing ? 'w-full' : 'max-w-[70%]'}
-            ${isEditing ? 'flex-grow' : 'flex-grow-0'}
-            ${isEditable ? '' : 'ml-auto'}`}
+          className={cn(
+            'px-5 py-2.5 relative bg-accent rounded-3xl',
+            inlineElements.length ? 'rounded-tr-lg' : '',
+            isEditing ? 'w-full flex-grow' : 'max-w-[70%] flex-grow-0',
+            isEditable ? '' : 'ml-auto'
+          )}
         >
           {isEditing ? (
-            <div className="bg-accent rounded-3xl flex flex-col">
+            <div className="bg-accent flex flex-col">
               <AutoResizeTextarea
                 id="edit-chat-input"
                 autoFocus

diff --git a/frontend/src/components/chat/Messages/Message/index.tsx b/frontend/src/components/chat/Messages/Message/index.tsx
@@ -84,9 +84,9 @@ const Message = memo(
               {/* User message is displayed differently */}
               {isUserMessage ? (
                 <div className="flex flex-col flex-grow max-w-full">
-                  <UserMessage message={message}>
+                  <UserMessage message={message} elements={elements}>
                     <MessageContent
-                      elements={elements}
+                      elements={[]}
                       message={message}
                       allowHtml={allowHtml}
                       latex={latex}

diff --git a/frontend/src/components/chat/MessagesContainer/index.tsx b/frontend/src/components/chat/MessagesContainer/index.tsx
@@ -147,6 +147,7 @@ const MessagesContainer = ({ navigate }: Props) => {
     <MessageContext.Provider value={memoizedContext}>
       <Messages
         indent={0}
+        isRunning={loading}
         messages={messages}
         elements={elements}
         actions={actions}