Add mention detection to update_message (#302)

* Add mention detection to update_message

  This patch extends the mention detection functionality from add_message
  (added in PR #235) to update_message. This enables automatic @mention
  detection when messages are updated, which is particularly useful for
  streaming scenarios where message content is built incrementally.

  Changes:
  - Refactored mention detection logic into reusable _extract_mentions() helper method
  - Updated add_message to use the new helper method
  - Added mention detection to update_message that scans the complete message body
  - Added three test cases covering mention updates, append with mentions, and deduplication
  - Fixed regex pattern to use raw string to avoid SyntaxWarning

  The implementation ensures no duplicate mentions are added when the same
  message is updated multiple times, making it safe for streaming use cases.

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

* changed variable name

* changed update_message with mentions tests to use sorted() instead of set() for better comparison

* revert test_add_message_includes_mentions to use set()

* changed to a set of actions

* updated doc string

* changed trigger actions API to be a list of callbacks, not strings, added a utils file for future actions

* updated name of callback to just be find_mentions

* fixed the test naming of callback

* fixed one last import name from find_mentions_callback to just find_mentions

---------

Co-authored-by: Abigayle-Mercer <[email protected]>
Co-authored-by: Zachary Sailer <[email protected]>

Author: 3 people

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

@@ -9,6 +9,7 @@
 from uuid import uuid4
 from ..models import message_asdict_factory, Message, NewMessage, User
 from ..ychat import YChat
+from ..utils import find_mentions
 
 USER = User(
     username=str(uuid4()),
@@ -167,6 +168,63 @@ def test_update_message_should_append_content():
     assert message_dict["sender"] == msg.sender
 
 
+def test_update_message_includes_mentions():
+    chat = YChat()
+    chat.set_user(USER)
+    chat.set_user(USER2)
+    chat.set_user(USER3)
+
+    new_msg = create_new_message(f"@{USER2.mention_name} Hello!")
+    msg_id = chat.add_message(new_msg)
+    msg = chat.get_message(msg_id)
+    assert msg
+    assert msg.mentions == [USER2.username]
+
+    msg.body = f"@{USER3.mention_name} Goodbye!"
+    chat.update_message(msg, trigger_actions=[find_mentions])
+    updated_msg = chat.get_message(msg_id)
+    assert updated_msg
+    assert updated_msg.mentions == [USER3.username]
+
+
+def test_update_message_append_includes_mentions():
+    chat = YChat()
+    chat.set_user(USER)
+    chat.set_user(USER2)
+    chat.set_user(USER3)
+
+    new_msg = create_new_message(f"@{USER2.mention_name} Hello!")
+    msg_id = chat.add_message(new_msg)
+    msg = chat.get_message(msg_id)
+    assert msg
+    assert msg.mentions == [USER2.username]
+
+    msg.body = f" and @{USER3.mention_name}!"
+    chat.update_message(msg, append=True, trigger_actions=[find_mentions])
+    updated_msg = chat.get_message(msg_id)
+    assert updated_msg
+    assert sorted(updated_msg.mentions) == sorted([USER2.username, USER3.username])
+
+
+def test_update_message_append_no_duplicate_mentions():
+    chat = YChat()
+    chat.set_user(USER)
+    chat.set_user(USER2)
+
+    new_msg = create_new_message(f"@{USER2.mention_name} Hello!")
+    msg_id = chat.add_message(new_msg)
+    msg = chat.get_message(msg_id)
+    assert msg
+    assert msg.mentions == [USER2.username]
+
+    msg.body = f" @{USER2.mention_name} again!"
+    chat.update_message(msg, append=True, trigger_actions=[find_mentions])
+    updated_msg = chat.get_message(msg_id)
+    assert updated_msg
+    assert updated_msg.mentions == [USER2.username]
+    assert len(updated_msg.mentions) == 1
+
+
 def test_indexes_by_id():
     chat = YChat()
     msg = create_new_message()

python/jupyterlab-chat/jupyterlab_chat/utils.py

@@ -0,0 +1,33 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+"""Utility functions for jupyter-chat."""
+
+import re
+from typing import TYPE_CHECKING, Set
+
+if TYPE_CHECKING:
+    from .models import Message
+    from .ychat import YChat
+
+
+def find_mentions(message: "Message", chat: "YChat") -> None:
+    """
+    Callback to extract and update mentions in a message.
+
+    Finds all @mentions in the message body and updates the message's mentions list
+    with the corresponding usernames.
+
+    Args:
+        message: The message object to update
+        chat: The YChat instance for accessing user data
+    """
+    mention_pattern = re.compile(r"@([\w-]+):?")
+    mentioned_names: Set[str] = set(re.findall(mention_pattern, message.body))
+    users = chat.get_users()
+    mentioned_usernames = []
+    for username, user in users.items():
+        if user.mention_name in mentioned_names and user.username not in mentioned_usernames:
+            mentioned_usernames.append(username)
+
+    message.mentions = mentioned_usernames

python/jupyterlab-chat/jupyterlab_chat/ychat.py

@@ -15,6 +15,7 @@
 import re
 
 from .models import message_asdict_factory, FileAttachment, NotebookAttachment, Message, NewMessage, User
+from .utils import find_mentions
 
 
 class YChat(YBaseDoc):
@@ -124,10 +125,18 @@ def _get_messages(self) -> list[dict]:
         """
         return self._ymessages.to_py() or []
 
-    def add_message(self, new_message: NewMessage) -> str:
+    def add_message(self, new_message: NewMessage, trigger_actions: list[Callable] | None = None) -> str:
         """
         Append a message to the document.
+
+        Args:
+            new_message: The message to add
+            trigger_actions: List of callbacks to execute on the message. Defaults to [find_mentions].
+                           Each callback receives (message, chat) as arguments.
         """
+        if trigger_actions is None:
+            trigger_actions = [find_mentions]
+
         timestamp: float = time.time()
         uid = str(uuid4())
         message = Message(
@@ -136,15 +145,9 @@ def add_message(self, new_message: NewMessage) -> str:
             id=uid,
         )
 
-        # find all mentioned users and add them as message mentions
-        mention_pattern = re.compile("@([\w-]+):?")
-        mentioned_names: Set[str] = set(re.findall(mention_pattern, message.body))
-        users = self.get_users()
-        mentioned_usernames = []
-        for username, user in users.items():
-            if user.mention_name in mentioned_names and user.username not in mentioned_usernames:
-                mentioned_usernames.append(username)
-        message.mentions = mentioned_usernames
+        # Execute all trigger action callbacks
+        for callback in trigger_actions:
+            callback(message, self)
 
         with self._ydoc.transaction():
             index = len(self._ymessages) - next((i for i, v in enumerate(self._get_messages()[::-1]) if v["time"] < timestamp), len(self._ymessages))
@@ -155,17 +158,27 @@ 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, trigger_actions: list[Callable] | None = None):
         """
         Update a message of the document.
-        If append is True, the content will be append to the previous content.
+
+        Args:
+            message: The message to update
+            append: If True, the content will be appended to the previous content
+            trigger_actions: List of callbacks to execute on the message. Each callback receives (message, chat) as arguments.
         """
         with self._ydoc.transaction():
             index = self._indexes_by_id[message.id]
             initial_message = self._ymessages[index]
             message.time = initial_message["time"]  # type:ignore[index]
             if append:
                 message.body = initial_message["body"] + message.body  # type:ignore[index]
+
+            # Execute all trigger action callbacks
+            if trigger_actions:
+                for callback in trigger_actions:
+                    callback(message, self)
+
             self._ymessages[index] = asdict(message, dict_factory=message_asdict_factory)
 
     def get_attachments(self) -> dict[str, Union[FileAttachment, NotebookAttachment]]: