diff --git a/src/frontend/package-lock.json b/src/frontend/package-lock.json
index 74e5994a7..5a0a618c5 100644
--- a/src/frontend/package-lock.json
+++ b/src/frontend/package-lock.json
@@ -6655,9 +6655,9 @@
       }
     },
     "node_modules/framer-motion": {
-      "version": "11.1.5",
-      "resolved": "https://registry.npmjs.org/framer-motion/-/framer-motion-11.1.5.tgz",
-      "integrity": "sha512-ogK5fc0GBUT3AjzMXPx7f74m5V1ByRqkKQARBVHpdkYTNDxb/WriANYD+5JBo1wklQQJ1HayDZtmofQLqZFcbw==",
+      "version": "11.1.7",
+      "resolved": "https://registry.npmjs.org/framer-motion/-/framer-motion-11.1.7.tgz",
+      "integrity": "sha512-cW11Pu53eDAXUEhv5hEiWuIXWhfkbV32PlgVISn7jRdcAiVrJ1S03YQQ0/DzoswGYYwKi4qYmHHjCzAH52eSdQ==",
       "dependencies": {
         "tslib": "^2.4.0"
       },
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx
index 93a0db3c3..39f754ff4 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx
@@ -235,7 +235,7 @@ dark:prose-invert"
                                       },
                                     ]}
                                     activeTab={"0"}
-                                    setActiveTab={() => {}}
+                                    setActiveTab={() => { }}
                                   />
                                 ) : (
                                   <code className={className} {...props}>
@@ -251,21 +251,6 @@ dark:prose-invert"
                       [chat.message, chatMessage]
                     )}
                   </div>
-                  {chat.files && (
-                    <div className="my-2 w-full">
-                      {chat.files.map((file, index) => {
-                        return (
-                          <div key={index} className="my-2 w-full">
-                            <FileCard
-                              fileName={"Generated File"}
-                              fileType={file.data_type}
-                              content={file.data}
-                            />
-                          </div>
-                        );
-                      })}
-                    </div>
-                  )}
                 </div>
               </div>
             </div>
@@ -292,33 +277,33 @@ dark:prose-invert"
                 <span className="prose text-primary word-break-break-word dark:prose-invert">
                   {promptOpen
                     ? template?.split("\n")?.map((line, index) => {
-                        const regex = /{([^}]+)}/g;
-                        let match;
-                        let parts: Array<JSX.Element | string> = [];
-                        let lastIndex = 0;
-                        while ((match = regex.exec(line)) !== null) {
-                          // Push text up to the match
-                          if (match.index !== lastIndex) {
-                            parts.push(line.substring(lastIndex, match.index));
-                          }
-                          // Push div with matched text
-                          if (chat.message[match[1]]) {
-                            parts.push(
-                              <span className="chat-message-highlight">
-                                {chat.message[match[1]]}
-                              </span>
-                            );
-                          }
+                      const regex = /{([^}]+)}/g;
+                      let match;
+                      let parts: Array<JSX.Element | string> = [];
+                      let lastIndex = 0;
+                      while ((match = regex.exec(line)) !== null) {
+                        // Push text up to the match
+                        if (match.index !== lastIndex) {
+                          parts.push(line.substring(lastIndex, match.index));
+                        }
+                        // Push div with matched text
+                        if (chat.message[match[1]]) {
+                          parts.push(
+                            <span className="chat-message-highlight">
+                              {chat.message[match[1]]}
+                            </span>
+                          );
+                        }
 
-                          // Update last index
-                          lastIndex = regex.lastIndex;
-                        }
-                        // Push text after the last match
-                        if (lastIndex !== line.length) {
-                          parts.push(line.substring(lastIndex));
-                        }
-                        return <p>{parts}</p>;
-                      })
+                        // Update last index
+                        lastIndex = regex.lastIndex;
+                      }
+                      // Push text after the last match
+                      if (lastIndex !== line.length) {
+                        parts.push(line.substring(lastIndex));
+                      }
+                      return <p>{parts}</p>;
+                    })
                     : chatMessage}
                 </span>
               </>
@@ -335,6 +320,21 @@ dark:prose-invert"
         )}
       </div>
       <div id={lastMessage ? "last-chat-message" : ""}></div>
+      {chat.files && (
+        <div className="my-2 w-full">
+          {chat.files.map((file, index) => {
+            return (
+              <div key={index} className="my-2 w-full">
+                <FileCard
+                  fileName={file.name}
+                  fileType={file.type}
+                  content={file.path}
+                />
+              </div>
+            );
+          })}
+        </div>
+      )}
     </>
   );
 }
diff --git a/src/frontend/src/modals/IOModal/components/chatView/fileComponent/index.tsx b/src/frontend/src/modals/IOModal/components/chatView/fileComponent/index.tsx
index 408e741c2..62b4a8071 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/fileComponent/index.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/fileComponent/index.tsx
@@ -2,6 +2,10 @@ import * as base64js from "base64-js";
 import { useState } from "react";
 import IconComponent from "../../../../../components/genericIconComponent";
 import { fileCardPropsType } from "../../../../../types/components";
+import useFlowsManagerStore from "../../../../../stores/flowsManagerStore";
+import { BACKEND_URL, BASE_URL_API } from "../../../../../constants/constants";
+
+const imgTypes = new Set(["png", "jpg"]);
 
 export default function FileCard({
   fileName,
@@ -9,18 +13,10 @@ export default function FileCard({
   fileType,
 }: fileCardPropsType): JSX.Element {
   const handleDownload = (): void => {
-    const byteArray = new Uint8Array(base64js.toByteArray(content));
-    const blob = new Blob([byteArray], { type: "application/octet-stream" });
-    const url = URL.createObjectURL(blob);
-    const link = document.createElement("a");
-    link.href = url;
-    link.download = fileName + ".png";
-    document.body.appendChild(link);
-    link.click();
-    document.body.removeChild(link);
-    URL.revokeObjectURL(url);
+    //TODO: update download function
   };
   const [isHovered, setIsHovered] = useState(false);
+  const currentFlowId = useFlowsManagerStore((state) => state.currentFlowId);
   function handleMouseEnter(): void {
     setIsHovered(true);
   }
@@ -28,7 +24,7 @@ export default function FileCard({
     setIsHovered(false);
   }
 
-  if (fileType === "image") {
+  if (imgTypes.has(fileType)) {
     return (
       <div
         className="relative h-1/4 w-1/4"
@@ -36,8 +32,8 @@ export default function FileCard({
         onMouseLeave={handleMouseLeave}
       >
         <img
-          src={`data:image/png;base64,${content}`}
-          alt="generated image"
+            src={`${BACKEND_URL.slice(0,BACKEND_URL.length-1)}${BASE_URL_API}files/images/${content}`}
+            alt="generated image"
           className="h-full  w-full rounded-lg"
         />
         {isHovered && (
@@ -60,10 +56,10 @@ export default function FileCard({
   return (
     <button onClick={handleDownload} className="file-card-modal-button">
       <div className="file-card-modal-div">
-        ooooooooooooooo{" "}
-        {fileType === "image" ? (
+        {" "}
+        {imgTypes.has(fileType) ? (
           <img
-            src={`data:image/png;base64,${content}`}
+            src={`${BACKEND_URL.slice(0,BACKEND_URL.length-1)}${BASE_URL_API}files/images/${content}`}
             alt=""
             className="h-8 w-8"
           />
diff --git a/src/frontend/src/modals/IOModal/components/chatView/index.tsx b/src/frontend/src/modals/IOModal/components/chatView/index.tsx
index cf27c3bfd..ff03e7e6a 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/index.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/index.tsx
@@ -67,7 +67,7 @@ export default function ChatView({
       .filter((output) => output.data.messages && output.data.messages.length > 0)
       .map((output, index) => {
         try {
-          const { sender, message, sender_name, stream_url } = output.data
+          const { sender, message, sender_name, stream_url,files } = output.data
             .messages[0] as ChatOutputType;
 
           const is_ai = sender === "Machine" || sender === null;
@@ -77,6 +77,8 @@ export default function ChatView({
             sender_name,
             componentId: output.id,
             stream_url: stream_url,
+            files
+            
           };
         } catch (e) {
           console.error(e);
diff --git a/src/frontend/src/types/chat/index.ts b/src/frontend/src/types/chat/index.ts
index 12c7a2a20..6fc4365c6 100644
--- a/src/frontend/src/types/chat/index.ts
+++ b/src/frontend/src/types/chat/index.ts
@@ -6,7 +6,7 @@ export type ChatMessageType = {
   template?: string;
   isSend: boolean;
   thought?: string;
-  files?: Array<{ data: string; type: string; data_type: string }>;
+  files?: Array<{ path: string; type: string; name: string }>;
   prompt?: string;
   chatKey?: string;
   componentId: string;