Fix CopyToCliboard functionality in UI (#411)

Co-authored-by: Philip Meier <[email protected]>
Quansight · May 13, 2024 · 0cb205f · 0cb205f
1 parent 6f7fbac
commit 0cb205f
Showing 1 changed file with 24 additions and 17 deletions.
diff --git a/ragna/deploy/_ui/central_view.py b/ragna/deploy/_ui/central_view.py
@@ -55,18 +55,38 @@ def __init__(
         on_click_source_info_callback: Optional[Callable] = None,
         timestamp=None,
         show_timestamp=True,
+        assistant_toolbar_visible=True,  # hide the toolbar during streaming
     ):
         css_class = f"message-content-{self.role}"
         self.content_pane = pn.pane.Markdown(
             content,
             css_classes=["message-content", css_class],
         )
 
+        # we make this available on the instance so that we can update the value later
+        self.clipboard_button = CopyToClipboardButton(
+            value=self.content_pane.object, title="Copy"
+        )
+
+        # we make this available on the instance so that we can toggle the visibility
+        self.assistant_toolbar = pn.Row(
+            self.clipboard_button,
+            pn.widgets.Button(
+                name="Source Info",
+                icon="info-circle",
+                css_classes=["source-info-button"],
+                on_click=lambda event: self.on_click_source_info_callback(
+                    event, self.sources
+                ),
+            ),
+            visible=assistant_toolbar_visible,
+        )
+
         if role == "assistant":
             assert sources is not None
             object = pn.Column(
                 self.content_pane,
-                self._copy_and_source_view_buttons(),
+                self.assistant_toolbar,
                 css_classes=["message-content-assistant-with-buttons"],
             )
         else:
@@ -92,22 +112,6 @@ def __init__(
         )
         self._stylesheets.append("css/chat_interface/chatmessage.css")
 
-    def _copy_and_source_view_buttons(self) -> pn.Row:
-        return pn.Row(
-            CopyToClipboardButton(
-                value=self.content_pane.object,
-                title="Copy",
-            ),
-            pn.widgets.Button(
-                name="Source Info",
-                icon="info-circle",
-                css_classes=["source-info-button"],
-                on_click=lambda event: self.on_click_source_info_callback(
-                    event, self.sources
-                ),
-            ),
-        )
-
     # This cannot be a bound method, because it creates a reference cycle when trying
     # to access the repr of the message. See
     # https://github.com/Quansight/ragna/issues/359 for details.
@@ -288,11 +292,14 @@ async def chat_callback(
                 user=self.get_user_from_role("assistant"),
                 sources=answer["sources"],
                 on_click_source_info_callback=self.on_click_source_info_wrapper,
+                assistant_toolbar_visible=False,
             )
             yield message
 
             async for chunk in answer_stream:
                 message.content_pane.object += chunk["content"]
+            message.clipboard_button.value = message.content_pane.object
+            message.assistant_toolbar.visible = True
 
         except Exception:
             yield RagnaChatMessage(