From 8448e5034cf287bc7fc091c1f02e8a2fc56037ac Mon Sep 17 00:00:00 2001
From: pablonyx <pablo@danswer.ai>
Date: Wed, 19 Feb 2025 12:42:36 -0800
Subject: [PATCH] k

---
 web/src/app/chat/ChatPage.tsx | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx
index 363eaf30936..9b682370743 100644
--- a/web/src/app/chat/ChatPage.tsx
+++ b/web/src/app/chat/ChatPage.tsx
@@ -2695,6 +2695,11 @@ export function ChatPage({
                                     ? messageHistory[i + 1]?.documents
                                     : undefined;
 
+                                const nextMessage =
+                                  messageHistory[i + 1]?.type === "assistant"
+                                    ? messageHistory[i + 1]
+                                    : undefined;
+
                                 return (
                                   <div
                                     className="text-text"
@@ -2723,7 +2728,10 @@ export function ChatPage({
                                             selectedMessageForDocDisplay ==
                                               secondLevelMessage?.messageId)
                                         }
-                                        isImprovement={message.isImprovement}
+                                        isImprovement={
+                                          message.isImprovement ||
+                                          nextMessage?.isImprovement
+                                        }
                                         secondLevelGenerating={
                                           (message.second_level_generating &&
                                             currentSessionChatState !==