added a is_done flag to make sure the full message is sent before mentioning personas

Author: abmercer

python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py

@@ -182,7 +182,7 @@ def test_update_message_includes_mentions():
 
     # Update the message to mention a different user
     msg.body = f"@{USER3.mention_name} Goodbye!"
-    chat.update_message(msg)
+    chat.update_message(msg, is_done=True)
     updated_msg = chat.get_message(msg_id)
     assert updated_msg
     assert set(updated_msg.mentions) == set([USER3.username])
@@ -203,7 +203,7 @@ def test_update_message_append_includes_mentions():
 
     # Append content with another mention
     msg.body = f" and @{USER3.mention_name}!"
-    chat.update_message(msg, append=True)
+    chat.update_message(msg, append=True, is_done=True)
     updated_msg = chat.get_message(msg_id)
     assert updated_msg
     # Should now mention both users
@@ -224,7 +224,7 @@ def test_update_message_append_no_duplicate_mentions():
 
     # Append content that mentions the same user again
     msg.body = f" @{USER2.mention_name} again!"
-    chat.update_message(msg, append=True)
+    chat.update_message(msg, append=True, is_done=True)
     updated_msg = chat.get_message(msg_id)
     assert updated_msg
     # Should only have one mention despite appearing twice

python/jupyterlab-chat/jupyterlab_chat/ychat.py

@@ -162,10 +162,11 @@ def add_message(self, new_message: NewMessage) -> str:
 
         return uid
 
-    def update_message(self, message: Message, append: bool = False):
+    def update_message(self, message: Message, append: bool = False, is_done: bool = False):
         """
         Update a message of the document.
-        If append is True, the content will be append to the previous content.
+        If append is True, the content will be appended to the previous content.
+        If is_done is True, mentions will be extracted and notifications triggered (use for streaming completion).
         """
         with self._ydoc.transaction():
             index = self._indexes_by_id[message.id]
@@ -175,7 +176,8 @@ def update_message(self, message: Message, append: bool = False):
                 message.body = initial_message["body"] + message.body  # type:ignore[index]
 
             # Extract and update mentions from the message body
-            message.mentions = self._extract_mentions(message.body)
+            if is_done:
+                message.mentions = self._extract_mentions(message.body)
 
             self._ymessages[index] = asdict(message, dict_factory=message_asdict_factory)