From cc051ceb4be756969a41b0431321b194d5c58490 Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Mon, 4 Dec 2023 20:40:54 -0500
Subject: [PATCH] Show generated images in chat interface on Web client

---
 src/khoj/interface/web/chat.html | 43 +++++++++++++++++++++++++++++---
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/src/khoj/interface/web/chat.html b/src/khoj/interface/web/chat.html
index a55d8e29..39cb6e77 100644
--- a/src/khoj/interface/web/chat.html
+++ b/src/khoj/interface/web/chat.html
@@ -183,12 +183,18 @@ To get started, just start typing below. You can also type / to see a list of co
                         referenceSection.appendChild(polishedReference);
                     }
                 }
-        }
+            }
 
             return numOnlineReferences;
-                    }
+        }
+
+        function renderMessageWithReference(message, by, context=null, dt=null, onlineContext=null, intentType=null) {
+            if (intentType === "text-to-image") {
+                let imageMarkdown = `![](${message})`;
+                renderMessage(imageMarkdown, by, dt);
+                return;
+            }
 
-        function renderMessageWithReference(message, by, context=null, dt=null, onlineContext=null) {
             if (context == null && onlineContext == null) {
                 renderMessage(message, by, dt);
                 return;
@@ -253,6 +259,26 @@ To get started, just start typing below. You can also type / to see a list of co
             // Remove any text between <s>[INST] and </s> tags. These are spurious instructions for the AI chat model.
             newHTML = newHTML.replace(/<s>\[INST\].+(<\/s>)?/g, '');
 
+            // Customize the rendering of images
+            md.renderer.rules.image = function(tokens, idx, options, env, self) {
+                let token = tokens[idx];
+
+                // Get image source url. Only render images with src links
+                let srcIndex = token.attrIndex('src');
+                if (srcIndex < 0) { return ''; }
+                let src = token.attrs[srcIndex][1];
+
+                // Wrap the image in a link
+                var aStart = `<a href="${src}" target="_blank">`;
+                var aEnd = '</a>';
+
+                // Add class="text-to-image" to images
+                token.attrPush(['class', 'text-to-image']);
+
+                // Use the default renderer to render image markdown format
+                return aStart + self.renderToken(tokens, idx, options) + aEnd;
+            };
+
             // Render markdown
             newHTML = md.render(newHTML);
             // Get any elements with a class that starts with "language"
@@ -414,6 +440,9 @@ To get started, just start typing below. You can also type / to see a list of co
                                 if (chunk.startsWith("{") && chunk.endsWith("}")) {
                                     try {
                                         const responseAsJson = JSON.parse(chunk);
+                                        if (responseAsJson.imageUrl) {
+                                            rawResponse += `![${query}](${responseAsJson.imageUrl})`;
+                                        }
                                         if (responseAsJson.detail) {
                                             rawResponse += responseAsJson.detail;
                                         }
@@ -516,7 +545,7 @@ To get started, just start typing below. You can also type / to see a list of co
                 .then(response => {
                     // Render conversation history, if any
                     response.forEach(chat_log => {
-                        renderMessageWithReference(chat_log.message, chat_log.by, chat_log.context, new Date(chat_log.created), chat_log.onlineContext);
+                        renderMessageWithReference(chat_log.message, chat_log.by, chat_log.context, new Date(chat_log.created), chat_log.onlineContext, chat_log.intent?.type);
                     });
                 })
                 .catch(err => {
@@ -902,6 +931,9 @@ To get started, just start typing below. You can also type / to see a list of co
             margin-top: -10px;
             transform: rotate(-60deg)
         }
+        img.text-to-image {
+            max-width: 60%;
+        }
 
         #chat-footer {
             padding: 0;
@@ -1029,6 +1061,9 @@ To get started, just start typing below. You can also type / to see a list of co
                 margin: 4px;
                 grid-template-columns: auto;
             }
+            img.text-to-image {
+                max-width: 100%;
+            }
         }
         @media only screen and (min-width: 700px) {
             body {