From 9daaae0fdbcd1417868870a232dd99aeec716921 Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Fri, 11 Oct 2024 01:11:34 -0700
Subject: [PATCH] Render inline any image files output by code in message

Update regex to also include any links to code generated images that
aren't explicitly meant to be displayed inline. This allows folks to
download the image (unlike the fake link that doesn't work created by
model)
---
 src/interface/web/app/common/chatFunctions.ts               | 6 +++---
 .../web/app/components/chatMessage/chatMessage.tsx          | 4 ++--
 src/khoj/routers/helpers.py                                 | 1 +
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/interface/web/app/common/chatFunctions.ts b/src/interface/web/app/common/chatFunctions.ts
index e6c402a3..6035bf78 100644
--- a/src/interface/web/app/common/chatFunctions.ts
+++ b/src/interface/web/app/common/chatFunctions.ts
@@ -116,7 +116,7 @@ export function processMessageChunk(
         if (context) currentMessage.context = context;
 
         // Replace file links with base64 data
-        currentMessage.rawResponse = replaceFileLinksWithBase64(
+        currentMessage.rawResponse = renderCodeGenImageInline(
             currentMessage.rawResponse,
             codeContext,
         );
@@ -186,12 +186,12 @@ export function handleImageResponse(imageJson: any, liveStream: boolean): Respon
     return reference;
 }
 
-export function replaceFileLinksWithBase64(message: string, codeContext: CodeContext) {
+export function renderCodeGenImageInline(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");
+            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);
diff --git a/src/interface/web/app/components/chatMessage/chatMessage.tsx b/src/interface/web/app/components/chatMessage/chatMessage.tsx
index 75c2685b..31aa4b48 100644
--- a/src/interface/web/app/components/chatMessage/chatMessage.tsx
+++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx
@@ -10,7 +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 { renderCodeGenImageInline } from "@/app/common/chatFunctions";
 
 import {
     ThumbsUp,
@@ -379,7 +379,7 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
         }
 
         // Replace file links with base64 data
-        message = replaceFileLinksWithBase64(message, props.chatMessage.codeContext);
+        message = renderCodeGenImageInline(message, props.chatMessage.codeContext);
 
         // Add code context files to the message
         if (props.chatMessage.codeContext) {
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index 3b97e694..4814d390 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -452,6 +452,7 @@ async def infer_webpage_urls(
     # Validate that the response is a non-empty, JSON-serializable list of URLs
     try:
         response = response.strip()
+        response = remove_json_codeblock(response)
         urls = json.loads(response)
         valid_unique_urls = {str(url).strip() for url in urls["links"] if is_valid_url(url)}
         if is_none_or_empty(valid_unique_urls):