diff --git a/src/interface/web/app/components/chatMessage/chatMessage.module.css b/src/interface/web/app/components/chatMessage/chatMessage.module.css
index b055d0a5..2abfdda7 100644
--- a/src/interface/web/app/components/chatMessage/chatMessage.module.css
+++ b/src/interface/web/app/components/chatMessage/chatMessage.module.css
@@ -77,6 +77,21 @@ div.imageWrapper img {
     border-radius: 8px;
 }
 
+div.khoj div.imageWrapper img {
+    height: 512px;
+}
+
+div.khoj div.imageWrapper {
+    flex: 1 1 auto;
+}
+
+div.khoj div.imagesContainer {
+    display: flex;
+    flex-wrap: wrap;
+    flex-direction: row;
+    overflow-x: hidden;
+}
+
 div.chatMessageContainer > img {
     width: auto;
     height: auto;
@@ -178,4 +193,9 @@ div.trainOfThoughtElement ul {
     div.youfullHistory {
         max-width: 90%;
     }
+
+    div.khoj div.imageWrapper img {
+        width: 100%;
+        height: auto;
+    }
 }
diff --git a/src/khoj/processor/conversation/utils.py b/src/khoj/processor/conversation/utils.py
index 57fc0f12..798522e6 100644
--- a/src/khoj/processor/conversation/utils.py
+++ b/src/khoj/processor/conversation/utils.py
@@ -293,7 +293,7 @@ def save_to_conversation_log(
             "turnId": turn_id,
             "images": generated_images,
             "queryFiles": [file.model_dump(mode="json") for file in raw_generated_files],
-            "excalidrawDiagram": str(generated_excalidraw_diagram),
+            "excalidrawDiagram": str(generated_excalidraw_diagram) if generated_excalidraw_diagram else None,
         },
         conversation_log=meta_log.get("chat", []),
     )