Render images files output by code interpreter in message on web app

src/interface/web/app/common/chatFunctions.ts

@@ -115,6 +115,33 @@ export function processMessageChunk(
         if (onlineContext) currentMessage.onlineContext = onlineContext;
         if (context) currentMessage.context = context;
 
+        // Replace file links with base64 data
+        currentMessage.rawResponse = replaceFileLinksWithBase64(
+            currentMessage.rawResponse,
+            codeContext,
+        );
+
+        // Add code context files to the message
+        if (codeContext) {
+            Object.entries(codeContext).forEach(([key, value]) => {
+                value.results.output_files?.forEach((file) => {
+                    if (file.filename.endsWith(".png") || file.filename.endsWith(".jpg")) {
+                        // Don't add the image again if it's already in the message!
+                        if (!currentMessage.rawResponse.includes(`![${file.filename}](`)) {
+                            currentMessage.rawResponse += `\n\n![${file.filename}](data:image/png;base64,${file.b64_data})`;
+                        }
+                    } else if (
+                        file.filename.endsWith(".txt") ||
+                        file.filename.endsWith(".org") ||
+                        file.filename.endsWith(".md")
+                    ) {
+                        const decodedText = atob(file.b64_data);
+                        currentMessage.rawResponse += `\n\n\`\`\`\n${decodedText}\n\`\`\``;
+                    }
+                });
+            });
+        }
+
         // Mark current message streaming as completed
         currentMessage.completed = true;
     }
@@ -159,6 +186,22 @@ export function handleImageResponse(imageJson: any, liveStream: boolean): Respon
     return reference;
 }
 
+export function replaceFileLinksWithBase64(message: string, codeContext: CodeContext) {
+    if (!codeContext) return message;
+
+    Object.values(codeContext).forEach((contextData) => {
+        contextData.results.output_files?.forEach((file) => {
+            const regex = new RegExp(`!\\[.*?\\]\\(.*${file.filename}\\)`, "g");
+            if (file.filename.match(/\.(png|jpg|jpeg|gif|webp)$/i)) {
+                const replacement = `![${file.filename}](data:image/${file.filename.split(".").pop()};base64,${file.b64_data})`;
+                message = message.replace(regex, replacement);
+            }
+        });
+    });
+
+    return message;
+}
+
 export function modifyFileFilterForConversation(
     conversationId: string | null,
     filenames: string[],

src/interface/web/app/components/chatMessage/chatMessage.tsx

@@ -10,6 +10,7 @@ import { createRoot } from "react-dom/client";
 import "katex/dist/katex.min.css";
 
 import { TeaserReferencesSection, constructAllReferences } from "../referencePanel/referencePanel";
+import { replaceFileLinksWithBase64 } from "@/app/common/chatFunctions";
 
 import {
     ThumbsUp,
@@ -377,6 +378,30 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
             message += `\n\n${props.chatMessage.intent["inferred-queries"][0]}`;
         }
 
+        // Replace file links with base64 data
+        message = replaceFileLinksWithBase64(message, props.chatMessage.codeContext);
+
+        // Add code context files to the message
+        if (props.chatMessage.codeContext) {
+            Object.entries(props.chatMessage.codeContext).forEach(([key, value]) => {
+                value.results.output_files?.forEach((file) => {
+                    if (file.filename.endsWith(".png") || file.filename.endsWith(".jpg")) {
+                        // Don't add the image again if it's already in the message!
+                        if (!message.includes(`![${file.filename}](`)) {
+                            message += `\n\n![${file.filename}](data:image/png;base64,${file.b64_data})`;
+                        }
+                    } else if (
+                        file.filename.endsWith(".txt") ||
+                        file.filename.endsWith(".org") ||
+                        file.filename.endsWith(".md")
+                    ) {
+                        const decodedText = atob(file.b64_data);
+                        message += `\n\n\`\`\`\n${decodedText}\n\`\`\``;
+                    }
+                });
+            });
+        }
+
         setTextRendered(message);
 
         // Render the markdown