diff --git a/src/interface/web/app/common/chatFunctions.ts b/src/interface/web/app/common/chatFunctions.ts
index 5e9ac321..aeb74ee8 100644
--- a/src/interface/web/app/common/chatFunctions.ts
+++ b/src/interface/web/app/common/chatFunctions.ts
@@ -139,33 +139,6 @@ export function processMessageChunk(
         if (onlineContext) currentMessage.onlineContext = onlineContext;
         if (context) currentMessage.context = context;
 
-        // Replace file links with base64 data
-        currentMessage.rawResponse = renderCodeGenImageInline(
-            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;
     }
@@ -200,10 +173,10 @@ export function renderCodeGenImageInline(message: string, codeContext: CodeConte
     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)) {
+            if (file.filename.match(/\.(png|jpg|jpeg)$/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)) {
+            } else if (file.filename.match(/\.(txt|org|md|csv|json)$/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 4869b5fa..a8fb2f6f 100644
--- a/src/interface/web/app/components/chatMessage/chatMessage.tsx
+++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx
@@ -421,6 +421,31 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
             }
         }
 
+        // Replace file links with base64 data
+        message = renderCodeGenImageInline(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") ||
+                        file.filename.endsWith(".csv") ||
+                        file.filename.endsWith(".json")
+                    ) {
+                        message += `\n\n## ${file.filename}\n\`\`\`\n${file.b64_data}\n\`\`\`\n`;
+                    }
+                });
+            });
+        }
+
         // Handle user attached images rendering
         let messageForClipboard = message;
         let messageToRender = message;
@@ -446,47 +471,6 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
             messageToRender = `${userImagesInHtml}${messageToRender}`;
         }
 
-        if (props.chatMessage.intent && props.chatMessage.intent.type == "text-to-image") {
-            message = `![generated image](data:image/png;base64,${message})`;
-        } else if (props.chatMessage.intent && props.chatMessage.intent.type == "text-to-image2") {
-            message = `![generated image](${message})`;
-        } else if (
-            props.chatMessage.intent &&
-            props.chatMessage.intent.type == "text-to-image-v3"
-        ) {
-            message = `![generated image](data:image/webp;base64,${message})`;
-        }
-        if (
-            props.chatMessage.intent &&
-            props.chatMessage.intent.type.includes("text-to-image") &&
-            props.chatMessage.intent["inferred-queries"]?.length > 0
-        ) {
-            message += `\n\n${props.chatMessage.intent["inferred-queries"][0]}`;
-        }
-
-        // Replace file links with base64 data
-        message = renderCodeGenImageInline(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")
-                    ) {
-                        message += `\n\n## ${file.filename}\n\`\`\`\n${file.b64_data}\n\`\`\`\n`;
-                    }
-                });
-            });
-        }
-
         // Set the message text
         setTextRendered(messageForClipboard);
 
diff --git a/src/interface/web/app/components/referencePanel/referencePanel.tsx b/src/interface/web/app/components/referencePanel/referencePanel.tsx
index 785d05d8..53016cfc 100644
--- a/src/interface/web/app/components/referencePanel/referencePanel.tsx
+++ b/src/interface/web/app/components/referencePanel/referencePanel.tsx
@@ -103,7 +103,7 @@ interface CodeContextReferenceCardProps {
 
 function CodeContextReferenceCard(props: CodeContextReferenceCardProps) {
     const fileIcon = getIconFromFilename(".py", "w-6 h-6 text-muted-foreground inline-flex mr-2");
-    const snippet = DOMPurify.sanitize(props.code);
+    const sanitizedCodeSnippet = DOMPurify.sanitize(props.code.replace(/\n/g, "<br/>"));
     const [isHovering, setIsHovering] = useState(false);
 
     return (
@@ -123,9 +123,8 @@ function CodeContextReferenceCard(props: CodeContextReferenceCardProps) {
                         </h3>
                         <p
                             className={`${props.showFullContent ? "block" : "overflow-hidden line-clamp-2"}`}
-                        >
-                            {snippet}
-                        </p>
+                            dangerouslySetInnerHTML={{ __html: sanitizedCodeSnippet }}
+                        ></p>
                     </Card>
                 </PopoverTrigger>
                 <PopoverContent className="w-[400px] mx-2">
@@ -136,7 +135,10 @@ function CodeContextReferenceCard(props: CodeContextReferenceCardProps) {
                             {fileIcon}
                             Code
                         </h3>
-                        <p className={`overflow-hidden line-clamp-3`}>{snippet}</p>
+                        <p
+                            className={`overflow-hidden line-clamp-5`}
+                            dangerouslySetInnerHTML={{ __html: sanitizedCodeSnippet }}
+                        ></p>
                     </Card>
                 </PopoverContent>
             </Popover>
diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py
index 33476a97..7c79809f 100644
--- a/src/khoj/routers/api_chat.py
+++ b/src/khoj/routers/api_chat.py
@@ -1131,6 +1131,7 @@ async def chat(
                 conversation_id=conversation_id,
                 compiled_references=compiled_references,
                 online_results=online_results,
+                code_results=code_results,
                 query_images=uploaded_images,
                 train_of_thought=train_of_thought,
                 attached_file_context=attached_file_context,
@@ -1192,6 +1193,7 @@ async def chat(
                 conversation_id=conversation_id,
                 compiled_references=compiled_references,
                 online_results=online_results,
+                code_results=code_results,
                 query_images=uploaded_images,
                 train_of_thought=train_of_thought,
                 attached_file_context=attached_file_context,