From af0215765ca59373f08c491e6f927d9f64dbb3c6 Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Fri, 8 Nov 2024 17:19:42 -0800
Subject: [PATCH] Decode code text output files from b64 to str to ease client
 processing

---
 src/interface/web/app/common/chatFunctions.ts            | 4 ++++
 .../web/app/components/chatMessage/chatMessage.tsx       | 4 ++--
 src/khoj/processor/tools/run_code.py                     | 9 +++++++++
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/interface/web/app/common/chatFunctions.ts b/src/interface/web/app/common/chatFunctions.ts
index b4777ba8..5e9ac321 100644
--- a/src/interface/web/app/common/chatFunctions.ts
+++ b/src/interface/web/app/common/chatFunctions.ts
@@ -203,6 +203,10 @@ export function renderCodeGenImageInline(message: string, codeContext: CodeConte
             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);
+            } else if (file.filename.match(/\.(txt|org|md)$/i)) {
+                // render output files generated by codegen as downloadable links
+                const replacement = `![${file.filename}](data:text/plain;base64,${file.b64_data})`;
+                message = message.replace(regex, replacement);
             }
         });
     });
diff --git a/src/interface/web/app/components/chatMessage/chatMessage.tsx b/src/interface/web/app/components/chatMessage/chatMessage.tsx
index aae271b1..4869b5fa 100644
--- a/src/interface/web/app/components/chatMessage/chatMessage.tsx
+++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx
@@ -420,6 +420,7 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
                 message += `\n\n${inferredQueries[0]}`;
             }
         }
+
         // Handle user attached images rendering
         let messageForClipboard = message;
         let messageToRender = message;
@@ -480,8 +481,7 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
                         file.filename.endsWith(".org") ||
                         file.filename.endsWith(".md")
                     ) {
-                        const decodedText = atob(file.b64_data);
-                        message += `\n\n\`\`\`\n${decodedText}\n\`\`\``;
+                        message += `\n\n## ${file.filename}\n\`\`\`\n${file.b64_data}\n\`\`\`\n`;
                     }
                 });
             });
diff --git a/src/khoj/processor/tools/run_code.py b/src/khoj/processor/tools/run_code.py
index fe89da1c..5366cd02 100644
--- a/src/khoj/processor/tools/run_code.py
+++ b/src/khoj/processor/tools/run_code.py
@@ -2,7 +2,9 @@ import base64
 import datetime
 import json
 import logging
+import mimetypes
 import os
+from pathlib import Path
 from typing import Any, Callable, List, NamedTuple, Optional
 
 import aiohttp
@@ -160,6 +162,13 @@ async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_ur
             if response.status == 200:
                 result: dict[str, Any] = await response.json()
                 result["code"] = cleaned_code
+                # Store decoded output files
+                for output_file in result.get("output_files", []):
+                    # Decode text files as UTF-8
+                    if mimetypes.guess_type(output_file["filename"])[0].startswith("text/") or Path(
+                        output_file["filename"]
+                    ).suffix in [".org", ".md", ".json"]:
+                        output_file["b64_data"] = base64.b64decode(output_file["b64_data"]).decode("utf-8")
                 return result
             else:
                 return {