From 491a1a949a9abddbd0a45cb84ba3e0365b362e04 Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Mon, 27 Nov 2023 11:45:36 -0800
Subject: [PATCH] Render chat responses as markdown in Desktop client too

---
 src/interface/desktop/chat.html  | 84 ++++++++++++++++++++++++++------
 src/khoj/interface/web/chat.html |  5 +-
 2 files changed, 74 insertions(+), 15 deletions(-)

diff --git a/src/interface/desktop/chat.html b/src/interface/desktop/chat.html
index efb505f4..599a8505 100644
--- a/src/interface/desktop/chat.html
+++ b/src/interface/desktop/chat.html
@@ -8,6 +8,7 @@
         <link rel="manifest" href="/static/khoj_chat.webmanifest">
         <link rel="stylesheet" href="./assets/khoj.css">
     </head>
+    <script type="text/javascript" src="./assets/markdown-it.min.js"></script>
     <script src="./utils.js"></script>
 
     <script>
@@ -120,8 +121,7 @@
             // Create a new div for the chat message text and append it to the chat message
             let chatMessageText = document.createElement('div');
             chatMessageText.className = `chat-message-text ${by}`;
-            let textNode = document.createTextNode(formattedMessage);
-            chatMessageText.appendChild(textNode);
+            chatMessageText.appendChild(formattedMessage);
             chatMessage.appendChild(chatMessageText);
 
             // Append annotations div to the chat message
@@ -230,16 +230,49 @@
         }
 
         function formatHTMLMessage(htmlMessage) {
-            // Replace any ``` with <div class="programmatic-output">
-            let newHTML = htmlMessage.replace(/```([\s\S]*?)```/g, '<div class="programmatic-output"><button class="copy-button" onclick="copyProgrammaticOutput(event)">Copy</button>$1</div>');
-            // Replace any ** with <b> and __ with <u>
-            newHTML = newHTML.replace(/\*\*([\s\S]*?)\*\*/g, '<b>$1</b>');
-            newHTML = newHTML.replace(/__([\s\S]*?)__/g, '<u>$1</u>');
+            var md = window.markdownit();
+            let newHTML = htmlMessage;
+
             // 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, '');
-            // For any text that has single backticks, replace them with <code> tags
-            newHTML = newHTML.replace(/`([^`]+)`/g, '<code class="chat-response">$1</code>');
-            return newHTML;
+
+            // Render markdown
+            newHTML = md.render(newHTML);
+            // Get any elements with a class that starts with "language"
+            let element = document.createElement('div');
+            element.innerHTML = newHTML;
+            let codeBlockElements = element.querySelectorAll('[class^="language-"]');
+            // For each element, add a parent div with the class "programmatic-output"
+            codeBlockElements.forEach((codeElement) => {
+                // Create the parent div
+                let parentDiv = document.createElement('div');
+                parentDiv.classList.add("programmatic-output");
+                // Add the parent div before the code element
+                codeElement.parentNode.insertBefore(parentDiv, codeElement);
+                // Move the code element into the parent div
+                parentDiv.appendChild(codeElement);
+                // Add a copy button to each element
+                let copyButton = document.createElement('button');
+                copyButton.classList.add("copy-button");
+                copyButton.innerHTML = "Copy";
+                copyButton.addEventListener('click', copyProgrammaticOutput);
+                codeElement.prepend(copyButton);
+            });
+
+            // Get all code elements that have no class.
+            let codeElements = element.querySelectorAll('code:not([class])');
+            codeElements.forEach((codeElement) => {
+                // Add the class "chat-response" to each element
+                codeElement.classList.add("chat-response");
+            });
+
+            let anchorElements = element.querySelectorAll('a');
+            anchorElements.forEach((anchorElement) => {
+                // Add the class "inline-chat-link" to each element
+                anchorElement.classList.add("inline-chat-link");
+            });
+
+            return element
         }
 
         async function chat() {
@@ -299,8 +332,11 @@
                         reader.read().then(({ done, value }) => {
                             if (done) {
                                 // Append any references after all the data has been streamed
-                                newResponseText.appendChild(references);
+                                if (references != null) {
+                                    newResponseText.appendChild(references);
+                                }
                                 document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
+                                document.getElementById("chat-input").removeAttribute("disabled");
                                 return;
                             }
 
@@ -321,7 +357,6 @@
                                 let referenceExpandButton = document.createElement('button');
                                 referenceExpandButton.classList.add("reference-expand-button");
 
-
                                 let referenceSection = document.createElement('div');
                                 referenceSection.classList.add("reference-section");
                                 referenceSection.classList.add("collapsed");
@@ -387,7 +422,6 @@
                         });
                     }
                     readStream();
-                    document.getElementById("chat-input").removeAttribute("disabled");
                 });
         }
 
@@ -935,7 +969,6 @@
             background: var(--primary-hover);
         }
 
-
         .option-enabled:focus {
             outline: none !important;
             border:1px solid #475569;
@@ -948,6 +981,25 @@
             border-bottom: 1px dotted #475569;
         }
 
+        a.reference-link {
+            color: var(--main-text-color);
+            border-bottom: 1px dotted var(--main-text-color);
+        }
+
+        button.copy-button {
+            display: block;
+            border-radius: 4px;
+            background-color: var(--background-color);
+        }
+        button.copy-button:hover {
+            background: #f5f5f5;
+            cursor: pointer;
+        }
+
+        pre {
+            text-wrap: unset;
+        }
+
         div.khoj-empty-container {
             padding: 0;
             margin: 0;
@@ -1021,6 +1073,10 @@
             text-align: center;
         }
 
+        p {
+            margin: 0;
+        }
+
         div.programmatic-output {
             background-color: #f5f5f5;
             border: 1px solid #ddd;
diff --git a/src/khoj/interface/web/chat.html b/src/khoj/interface/web/chat.html
index 14d587f1..0cbd0f67 100644
--- a/src/khoj/interface/web/chat.html
+++ b/src/khoj/interface/web/chat.html
@@ -975,12 +975,15 @@ To get started, just start typing below. You can also type / to see a list of co
             border-radius: 4px;
             background-color: var(--background-color);
         }
-
         button.copy-button:hover {
             background: #f5f5f5;
             cursor: pointer;
         }
 
+        pre {
+            text-wrap: unset;
+        }
+
         @media (pointer: coarse), (hover: none) {
             abbr[title] {
                 position: relative;