From fb2393a940fd433b12cb406f93ce8f1c7ec7e475 Mon Sep 17 00:00:00 2001 From: Abigayle-Mercer Date: Thu, 30 Oct 2025 12:40:59 -0700 Subject: [PATCH 1/5] 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. --- .../jupyterlab_chat/tests/test_ychat.py | 64 +++++++++++++++++++ .../jupyterlab-chat/jupyterlab_chat/ychat.py | 27 +++++--- 2 files changed, 83 insertions(+), 8 deletions(-) diff --git a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py b/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py index b812549e..5dc1e724 100644 --- a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py +++ b/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py @@ -167,6 +167,70 @@ 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) + + # Add a message with one mention + 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 set(msg.mentions) == set([USER2.username]) + + # Update the message to mention a different user + msg.body = f"@{USER3.mention_name} Goodbye!" + chat.update_message(msg) + updated_msg = chat.get_message(msg_id) + assert updated_msg + assert set(updated_msg.mentions) == set([USER3.username]) + + +def test_update_message_append_includes_mentions(): + chat = YChat() + chat.set_user(USER) + chat.set_user(USER2) + chat.set_user(USER3) + + # Add a message with one mention + 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 set(msg.mentions) == set([USER2.username]) + + # Append content with another mention + msg.body = f" and @{USER3.mention_name}!" + chat.update_message(msg, append=True) + updated_msg = chat.get_message(msg_id) + assert updated_msg + # Should now mention both users + assert set(updated_msg.mentions) == set([USER2.username, USER3.username]) + + +def test_update_message_append_no_duplicate_mentions(): + chat = YChat() + chat.set_user(USER) + chat.set_user(USER2) + + # Add a message with a mention + 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 set(msg.mentions) == set([USER2.username]) + + # Append content that mentions the same user again + msg.body = f" @{USER2.mention_name} again!" + chat.update_message(msg, append=True) + updated_msg = chat.get_message(msg_id) + assert updated_msg + # Should only have one mention despite appearing twice + assert set(updated_msg.mentions) == set([USER2.username]) + + def test_indexes_by_id(): chat = YChat() msg = create_new_message() diff --git a/python/jupyterlab-chat/jupyterlab_chat/ychat.py b/python/jupyterlab-chat/jupyterlab_chat/ychat.py index a89e4fe0..1c1545cc 100644 --- a/python/jupyterlab-chat/jupyterlab_chat/ychat.py +++ b/python/jupyterlab-chat/jupyterlab_chat/ychat.py @@ -118,6 +118,20 @@ def get_messages(self) -> list[Message]: message_dicts = self._get_messages() return [Message(**message_dict) for message_dict in message_dicts] + def _extract_mentions(self, body: str) -> list[str]: + """ + Extract mentioned usernames from a message body. + Finds all @mentions in the body and returns the corresponding usernames. + """ + mention_pattern = re.compile(r"@([\w-]+):?") + mentioned_names: Set[str] = set(re.findall(mention_pattern, 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) + return mentioned_usernames + def _get_messages(self) -> list[dict]: """ Returns the messages of the document as dict. @@ -137,14 +151,7 @@ def add_message(self, new_message: NewMessage) -> str: ) # 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 + message.mentions = self._extract_mentions(message.body) 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)) @@ -166,6 +173,10 @@ def update_message(self, message: Message, append: bool = False): message.time = initial_message["time"] # type:ignore[index] if append: 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) + self._ymessages[index] = asdict(message, dict_factory=message_asdict_factory) def get_attachments(self) -> dict[str, Union[FileAttachment, NotebookAttachment]]: From 2e81dcb0957e7a41361c8377b5ab06d370ded612 Mon Sep 17 00:00:00 2001 From: Abigayle-Mercer Date: Fri, 31 Oct 2025 09:53:59 -0700 Subject: [PATCH 2/5] added a is_done flag to make sure the full message is sent before mentioning personas --- .../jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py | 6 +++--- python/jupyterlab-chat/jupyterlab_chat/ychat.py | 8 +++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py b/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py index 5dc1e724..1d734cbb 100644 --- a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py +++ b/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 diff --git a/python/jupyterlab-chat/jupyterlab_chat/ychat.py b/python/jupyterlab-chat/jupyterlab_chat/ychat.py index 1c1545cc..093a3248 100644 --- a/python/jupyterlab-chat/jupyterlab_chat/ychat.py +++ b/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) From ac993e9e62d8c8df1b308b6f9c23544454b10335 Mon Sep 17 00:00:00 2001 From: Abigayle-Mercer Date: Sat, 1 Nov 2025 16:15:38 -0700 Subject: [PATCH 3/5] changed variable name --- .../jupyterlab_chat/tests/test_ychat.py | 6 +++--- python/jupyterlab-chat/jupyterlab_chat/ychat.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py b/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py index 1d734cbb..8e8ee0e2 100644 --- a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py +++ b/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, is_done=True) + chat.update_message(msg, find_mentions=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, is_done=True) + chat.update_message(msg, append=True, find_mentions=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, is_done=True) + chat.update_message(msg, append=True, find_mentions=True) updated_msg = chat.get_message(msg_id) assert updated_msg # Should only have one mention despite appearing twice diff --git a/python/jupyterlab-chat/jupyterlab_chat/ychat.py b/python/jupyterlab-chat/jupyterlab_chat/ychat.py index 093a3248..55ffc0a7 100644 --- a/python/jupyterlab-chat/jupyterlab_chat/ychat.py +++ b/python/jupyterlab-chat/jupyterlab_chat/ychat.py @@ -118,7 +118,7 @@ def get_messages(self) -> list[Message]: message_dicts = self._get_messages() return [Message(**message_dict) for message_dict in message_dicts] - def _extract_mentions(self, body: str) -> list[str]: + def _find_mentions(self, body: str) -> list[str]: """ Extract mentioned usernames from a message body. Finds all @mentions in the body and returns the corresponding usernames. @@ -151,7 +151,7 @@ def add_message(self, new_message: NewMessage) -> str: ) # find all mentioned users and add them as message mentions - message.mentions = self._extract_mentions(message.body) + message.mentions = self._find_mentions(message.body) 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)) @@ -162,11 +162,11 @@ def add_message(self, new_message: NewMessage) -> str: return uid - def update_message(self, message: Message, append: bool = False, is_done: bool = False): + def update_message(self, message: Message, append: bool = False, find_mentions: bool = False): """ Update a message of the document. 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). + If find_mentions is True, mentions will be extracted and notifications triggered (use for streaming completion). """ with self._ydoc.transaction(): index = self._indexes_by_id[message.id] @@ -176,8 +176,8 @@ def update_message(self, message: Message, append: bool = False, is_done: bool = message.body = initial_message["body"] + message.body # type:ignore[index] # Extract and update mentions from the message body - if is_done: - message.mentions = self._extract_mentions(message.body) + if find_mentions: + message.mentions = self._find_mentions(message.body) self._ymessages[index] = asdict(message, dict_factory=message_asdict_factory) From f6440ea1bee8ef638fe996e83d39484b1df56883 Mon Sep 17 00:00:00 2001 From: Abigayle-Mercer Date: Tue, 11 Nov 2025 12:41:10 -0800 Subject: [PATCH 4/5] changed update_message with mentions tests to use sorted() instead of set() for better comparison --- .../jupyterlab_chat/tests/test_ychat.py | 23 +++++++------------ 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py b/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py index 8e8ee0e2..ce9fcc95 100644 --- a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py +++ b/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py @@ -115,7 +115,7 @@ def test_add_message_includes_mentions(): msg = chat.get_message(msg_id) assert msg - assert set(msg.mentions) == set([USER2.username, USER3.username]) + assert sorted(msg.mentions) == sorted([USER2.username, USER3.username]) def test_get_message_should_return_the_message(): @@ -173,19 +173,17 @@ def test_update_message_includes_mentions(): chat.set_user(USER2) chat.set_user(USER3) - # Add a message with one mention 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 set(msg.mentions) == set([USER2.username]) + assert msg.mentions == [USER2.username] - # Update the message to mention a different user msg.body = f"@{USER3.mention_name} Goodbye!" chat.update_message(msg, find_mentions=True) updated_msg = chat.get_message(msg_id) assert updated_msg - assert set(updated_msg.mentions) == set([USER3.username]) + assert updated_msg.mentions == [USER3.username] def test_update_message_append_includes_mentions(): @@ -194,20 +192,17 @@ def test_update_message_append_includes_mentions(): chat.set_user(USER2) chat.set_user(USER3) - # Add a message with one mention 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 set(msg.mentions) == set([USER2.username]) + assert msg.mentions == [USER2.username] - # Append content with another mention msg.body = f" and @{USER3.mention_name}!" chat.update_message(msg, append=True, find_mentions=True) updated_msg = chat.get_message(msg_id) assert updated_msg - # Should now mention both users - assert set(updated_msg.mentions) == set([USER2.username, USER3.username]) + assert sorted(updated_msg.mentions) == sorted([USER2.username, USER3.username]) def test_update_message_append_no_duplicate_mentions(): @@ -215,20 +210,18 @@ def test_update_message_append_no_duplicate_mentions(): chat.set_user(USER) chat.set_user(USER2) - # Add a message with a mention 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 set(msg.mentions) == set([USER2.username]) + assert msg.mentions == [USER2.username] - # Append content that mentions the same user again msg.body = f" @{USER2.mention_name} again!" chat.update_message(msg, append=True, find_mentions=True) updated_msg = chat.get_message(msg_id) assert updated_msg - # Should only have one mention despite appearing twice - assert set(updated_msg.mentions) == set([USER2.username]) + assert updated_msg.mentions == [USER2.username] + assert len(updated_msg.mentions) == 1 def test_indexes_by_id(): From a4617406206e7a8ac596708857c6c762ab2868cd Mon Sep 17 00:00:00 2001 From: Abigayle-Mercer Date: Tue, 11 Nov 2025 13:01:11 -0800 Subject: [PATCH 5/5] revert test_add_message_includes_mentions to use set() --- python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py b/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py index ce9fcc95..53b4c0d5 100644 --- a/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py +++ b/python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py @@ -115,7 +115,7 @@ def test_add_message_includes_mentions(): msg = chat.get_message(msg_id) assert msg - assert sorted(msg.mentions) == sorted([USER2.username, USER3.username]) + assert set(msg.mentions) == set([USER2.username, USER3.username]) def test_get_message_should_return_the_message():