From 7de6081ef6c12fec709fb2a0f0cb585b6bbbdf15 Mon Sep 17 00:00:00 2001
From: Willy Douhard <willy.douhard@gmail.com>
Date: Fri, 13 Sep 2024 13:46:39 +0200
Subject: [PATCH] Refine feedback UI and improve type handling (#1325)

* Enhance type checking for output logging in callbacks
* Display feedback buttons at end of run regardless of last step type
* Prevent feedback UI when no data layer is configured
* Update imports to include useConfig from chainlit/react-client
---
 backend/chainlit/langchain/callbacks.py               |  5 +++--
 .../src/components/molecules/messages/Message.tsx     |  9 ++++++---
 .../src/components/molecules/messages/Messages.tsx    | 10 +++++++++-
 .../molecules/messages/components/FeedbackButtons.tsx | 11 ++++++++++-
 4 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/backend/chainlit/langchain/callbacks.py b/backend/chainlit/langchain/callbacks.py
index 8243018e89..b7e29105cc 100644
--- a/backend/chainlit/langchain/callbacks.py
+++ b/backend/chainlit/langchain/callbacks.py
@@ -587,12 +587,13 @@ def _on_run_update(self, run: Run) -> None:
         outputs = run.outputs or {}
         output_keys = list(outputs.keys())
         output = outputs
+
         if output_keys:
             output = outputs.get(output_keys[0], outputs)
-
+            
         if current_step:
             current_step.output = (
-                output[0] if isinstance(output, Sequence) and len(output) else output
+                output[0] if isinstance(output, Sequence) and not isinstance(output, str) and len(output) else output
             )
             current_step.end = utc_now()
             self._run_sync(current_step.update())
diff --git a/frontend/src/components/molecules/messages/Message.tsx b/frontend/src/components/molecules/messages/Message.tsx
index f85d49cb95..ac20f42676 100644
--- a/frontend/src/components/molecules/messages/Message.tsx
+++ b/frontend/src/components/molecules/messages/Message.tsx
@@ -186,9 +186,12 @@ const Message = memo(
                       {actions?.length ? (
                         <MessageActions message={message} actions={actions} />
                       ) : null}
-                      {scorableRun && isScorable ? (
-                        <MessageButtons message={message} run={scorableRun} />
-                      ) : null}
+                      <MessageButtons
+                        message={message}
+                        run={
+                          scorableRun && isScorable ? scorableRun : undefined
+                        }
+                      />
                     </Stack>
                   )}
                 </Stack>
diff --git a/frontend/src/components/molecules/messages/Messages.tsx b/frontend/src/components/molecules/messages/Messages.tsx
index b1923ed9c9..024874299d 100644
--- a/frontend/src/components/molecules/messages/Messages.tsx
+++ b/frontend/src/components/molecules/messages/Messages.tsx
@@ -79,11 +79,19 @@ const Messages = memo(
             // Score the current run
             const _scorableRun = m.type === 'run' ? m : scorableRun;
             // The message is scorable if it is the last assistant message of the run
-            const isScorable =
+
+            const isRunLastAssistantMessage =
               m ===
               _scorableRun?.steps?.findLast(
                 (_m) => _m.type === 'assistant_message'
               );
+
+            const isLastAssistantMessage =
+              messages.findLast((_m) => _m.type === 'assistant_message') === m;
+
+            const isScorable =
+              isRunLastAssistantMessage || isLastAssistantMessage;
+
             return (
               <Message
                 message={m}
diff --git a/frontend/src/components/molecules/messages/components/FeedbackButtons.tsx b/frontend/src/components/molecules/messages/components/FeedbackButtons.tsx
index d9ccc16f7f..f6d0dd374f 100644
--- a/frontend/src/components/molecules/messages/components/FeedbackButtons.tsx
+++ b/frontend/src/components/molecules/messages/components/FeedbackButtons.tsx
@@ -7,7 +7,11 @@ import IconButton from '@mui/material/IconButton';
 import Stack from '@mui/material/Stack';
 import Tooltip from '@mui/material/Tooltip';
 
-import { firstUserInteraction, useChatSession } from '@chainlit/react-client';
+import {
+  firstUserInteraction,
+  useChatSession,
+  useConfig
+} from '@chainlit/react-client';
 
 import Dialog from 'components/atoms/Dialog';
 import { AccentButton } from 'components/atoms/buttons/AccentButton';
@@ -30,6 +34,7 @@ interface Props {
 }
 
 const FeedbackButtons = ({ message }: Props) => {
+  const config = useConfig();
   const { onFeedbackUpdated, onFeedbackDeleted } = useContext(MessageContext);
   const [showFeedbackDialog, setShowFeedbackDialog] = useState<number>();
   const [commentInput, setCommentInput] = useState<string>();
@@ -39,6 +44,10 @@ const FeedbackButtons = ({ message }: Props) => {
   const [feedback, setFeedback] = useState(message.feedback?.value);
   const [comment, setComment] = useState(message.feedback?.comment);
 
+  if (!config.config?.dataPersistence) {
+    return null;
+  }
+
   const DownIcon = feedback === 0 ? ThumbDownFilledIcon : ThumbDownIcon;
   const UpIcon = feedback === 1 ? ThumbUpFilledIcon : ThumbUpIcon;