From f398e1eb0c4b1a349a471485c7d8bd096a6cb53b Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Fri, 10 Jan 2025 21:46:39 -0800
Subject: [PATCH] Add codeblock rendering for the mermaidjs diagram in obsidian

---
 src/interface/obsidian/src/chat_view.ts | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/interface/obsidian/src/chat_view.ts b/src/interface/obsidian/src/chat_view.ts
index 86d5e1f9..5097dc2d 100644
--- a/src/interface/obsidian/src/chat_view.ts
+++ b/src/interface/obsidian/src/chat_view.ts
@@ -487,7 +487,8 @@ export class KhojChatView extends KhojPaneView {
         inferredQueries?: string[],
         conversationId?: string,
         images?: string[],
-        excalidrawDiagram?: string
+        excalidrawDiagram?: string,
+        mermaidjsDiagram?: string
     ) {
         if (!message) return;
 
@@ -496,8 +497,9 @@ export class KhojChatView extends KhojPaneView {
             intentType?.includes("text-to-image") ||
             intentType === "excalidraw" ||
             (images && images.length > 0) ||
+            mermaidjsDiagram ||
             excalidrawDiagram) {
-            let imageMarkdown = this.generateImageMarkdown(message, intentType ?? "", inferredQueries, conversationId, images, excalidrawDiagram);
+            let imageMarkdown = this.generateImageMarkdown(message, intentType ?? "", inferredQueries, conversationId, images, excalidrawDiagram, mermaidjsDiagram);
             chatMessageEl = this.renderMessage(chatEl, imageMarkdown, sender, dt);
         } else {
             chatMessageEl = this.renderMessage(chatEl, message, sender, dt);
@@ -517,7 +519,7 @@ export class KhojChatView extends KhojPaneView {
         chatMessageBodyEl.appendChild(this.createReferenceSection(references));
     }
 
-    generateImageMarkdown(message: string, intentType: string, inferredQueries?: string[], conversationId?: string, images?: string[], excalidrawDiagram?: string): string {
+    generateImageMarkdown(message: string, intentType: string, inferredQueries?: string[], conversationId?: string, images?: string[], excalidrawDiagram?: string, mermaidjsDiagram?: string): string {
         let imageMarkdown = "";
         if (intentType === "text-to-image") {
             imageMarkdown = `![](data:image/png;base64,${message})`;
@@ -529,6 +531,8 @@ export class KhojChatView extends KhojPaneView {
             const domain = this.setting.khojUrl.endsWith("/") ? this.setting.khojUrl : `${this.setting.khojUrl}/`;
             const redirectMessage = `Hey, I'm not ready to show you diagrams yet here. But you can view it in ${domain}chat?conversationId=${conversationId}`;
             imageMarkdown = redirectMessage;
+        } else if (mermaidjsDiagram) {
+            imageMarkdown = "```mermaid\n" + mermaidjsDiagram + "\n```";
         } else if (images && images.length > 0) {
             for (let image of images) {
                 if (image.startsWith("https://")) {
@@ -908,6 +912,7 @@ export class KhojChatView extends KhojPaneView {
                         chatBodyEl.dataset.conversationId ?? "",
                         chatLog.images,
                         chatLog.excalidrawDiagram,
+                        chatLog.mermaidjsDiagram,
                     );
                     // push the user messages to the chat history
                     if (chatLog.by === "you") {
@@ -1012,7 +1017,7 @@ export class KhojChatView extends KhojPaneView {
     }
 
     handleJsonResponse(jsonData: any): void {
-        if (jsonData.image || jsonData.detail || jsonData.images || jsonData.excalidrawDiagram) {
+        if (jsonData.image || jsonData.detail || jsonData.images || jsonData.mermaidjsDiagram) {
             this.chatMessageState.rawResponse = this.handleImageResponse(jsonData, this.chatMessageState.rawResponse);
         } else if (jsonData.response) {
             this.chatMessageState.rawResponse = jsonData.response;
@@ -1407,6 +1412,8 @@ export class KhojChatView extends KhojPaneView {
             const domain = this.setting.khojUrl.endsWith("/") ? this.setting.khojUrl : `${this.setting.khojUrl}/`;
             const redirectMessage = `Hey, I'm not ready to show you diagrams yet here. But you can view it in ${domain}`;
             rawResponse += redirectMessage;
+        } else if (imageJson.mermaidjsDiagram) {
+            rawResponse += imageJson.mermaidjsDiagram;
         }
 
         // If response has detail field, response is an error message.