Periodically update the persona awareness to keep it alive (#1358)

* Periodically update the personna awareness state to keep it alive

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Stop updating personna awareness if there is no user in document awareness

* Improve warning message

* Fix typing in warning message

* Update jupyter chat dependency to >=0.11.0

* Change the BOT contant from dict to User (should not be used anyway with the personna)

* Fix typing for python 3.9

* Update awareness state every 0.8 * outdated_timeout, and never set _heartbeat_task to None

Co-authored-by: David L. Qiu <[email protected]>

* Clean unused imports

* remove awareness observer

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: David L. Qiu <[email protected]>
Co-authored-by: David L. Qiu <[email protected]>

Author: 3 people

packages/jupyter-ai/jupyter_ai/chat_handlers/base.py

@@ -275,11 +275,11 @@ def reply(self, body: str, _human_message=None) -> str:
         TODO: Either properly store & use reply state in YChat, or remove the
         `human_message` argument here.
         """
-        bot = self.ychat.get_user(BOT["username"])
+        bot = self.ychat.get_user(BOT.username)
         if not bot:
-            self.ychat.set_user(User(**BOT))
+            self.ychat.set_user(BOT)
 
-        id = self.ychat.add_message(NewMessage(body=body, sender=BOT["username"]))
+        id = self.ychat.add_message(NewMessage(body=body, sender=BOT.username))
         return id
 
     @property

packages/jupyter-ai/jupyter_ai/chat_handlers/utils/streaming.py

@@ -38,9 +38,9 @@ def __init__(self, ychat: YChat):
         self._stream_id: Optional[str] = None
 
     def _set_user(self):
-        bot = self.ychat.get_user(BOT["username"])
+        bot = self.ychat.get_user(BOT.username)
         if not bot:
-            self.ychat.set_user(User(**BOT))
+            self.ychat.set_user(BOT)
 
     def open(self):
         self._set_user()
@@ -60,7 +60,7 @@ def write(self, chunk: str) -> str:
         if not self._stream_id:
             self._set_user()
             self._stream_id = self.ychat.add_message(
-                NewMessage(body="", sender=BOT["username"])
+                NewMessage(body="", sender=BOT.username)
             )
 
         self._set_user()
@@ -69,7 +69,7 @@ def write(self, chunk: str) -> str:
                 id=self._stream_id,
                 body=chunk,
                 time=time.time(),
-                sender=BOT["username"],
+                sender=BOT.username,
                 raw_time=False,
             ),
             append=True,

packages/jupyter-ai/jupyter_ai/constants.py

@@ -1,8 +1,11 @@
+from jupyterlab_chat.models import User
+
 # The BOT currently has a fixed username, because this username is used has key in chats,
 # it needs to constant. Do we need to change it ?
-BOT = {
-    "username": "5f6a7570-7974-6572-6e61-75742d626f74",
-    "name": "Jupyternaut",
-    "display_name": "Jupyternaut",
-    "initials": "J",
-}
+BOT = User(
+    username="5f6a7570-7974-6572-6e61-75742d626f74",
+    name="Jupyternaut",
+    display_name="Jupyternaut",
+    initials="J",
+    bot=True,
+)

packages/jupyter-ai/jupyter_ai/personas/base_persona.py

@@ -199,6 +199,7 @@ def as_user(self) -> User:
             name=self.name,
             display_name=self.name,
             avatar_url=self.avatar_path,
+            bot=True,
         )
 
     def as_user_dict(self) -> dict[str, Any]:

packages/jupyter-ai/jupyter_ai/personas/persona_awareness.py

@@ -1,3 +1,4 @@
+import asyncio
 import random
 from contextlib import contextmanager
 from dataclasses import asdict
@@ -21,7 +22,7 @@ class PersonaAwareness:
 
     - This class optionally accepts a `User` object in the constructor. When
     passed, this class will automatically register this user in the awareness
-    dictionary.
+    dictionary on init.
 
     - This class works by manually setting `ydoc.awareness.client_id` before &
     after each method call. This class provides a `self.as_custom_client()`
@@ -38,6 +39,7 @@ class PersonaAwareness:
 
     _original_client_id: int
     _custom_client_id: int
+    _heartbeat_task: asyncio.Task
 
     def __init__(self, *, ychat: YChat, log: Logger, user: Optional[User]):
         # Bind instance attributes
@@ -54,12 +56,14 @@ def __init__(self, *, ychat: YChat, log: Logger, user: Optional[User]):
         self._original_client_id = self.awareness.client_id
         self._custom_client_id = random.getrandbits(32)
 
-        with self.as_custom_client():
-            self.awareness.set_local_state({})
-
+        # Initialize local awareness state using the custom client ID
+        self.set_local_state({})
         if self.user:
             self._register_user()
 
+        # Start the awareness heartbeat task
+        self._heartbeat_task = asyncio.create_task(self._start_heartbeat())
+
     @contextmanager
     def as_custom_client(self) -> "Iterator[None]":
         """
@@ -77,6 +81,17 @@ def as_custom_client(self) -> "Iterator[None]":
         finally:
             self.awareness.client_id = self._original_client_id
 
+    @property
+    def outdated_timeout(self) -> int:
+        """
+        Returns the outdated timeout of the document awareness, in milliseconds.
+        The timeout value should be 30000 milliseconds (30 seconds), according to the
+        default value in `y-protocols.awareness` and `pycrdt.Awareness`.
+        - https://github.com/yjs/y-protocols/blob/2d8cd5c06b3925fbf9b5215dc341f8096a0a8d5c/awareness.js#L13
+        - https://github.com/y-crdt/pycrdt/blob/e269a3e63ad7986a3349e2d2bc7bd5f0dfca9c79/python/pycrdt/_awareness.py#L23
+        """
+        return self.awareness._outdated_timeout
+
     def _register_user(self):
         if not self.user:
             return
@@ -85,9 +100,43 @@ def _register_user(self):
             self.awareness.set_local_state_field("user", asdict(self.user))
 
     def get_local_state(self) -> Optional[dict[str, Any]]:
+        """
+        Returns the local state of the awareness instance.
+        """
         with self.as_custom_client():
             return self.awareness.get_local_state()
 
+    def set_local_state(self, state: dict[str, Any]) -> None:
+        """
+        Sets the local state of the awareness instance to the provided state.
+        This method is used to update the local state of the awareness instance
+        with a new state dictionary.
+        """
+        with self.as_custom_client():
+            self.awareness.set_local_state(state)
+
     def set_local_state_field(self, field: str, value: Any) -> None:
+        """
+        Sets a specific field in the local state of the awareness instance.
+        """
         with self.as_custom_client():
             self.awareness.set_local_state_field(field, value)
+
+    async def _start_heartbeat(self):
+        """
+        Background task that updates this instance's local state every
+        `0.8 * self.outdated_timeout` milliseconds. `pycrdt` and `yjs`
+        automatically disconnect clients if they do not make updates in
+        a long time (default: 30000 ms). This task keeps personas alive
+        after 30 seconds of no usage in each chat session.
+        """
+        while True:
+            await asyncio.sleep(0.8 * self.outdated_timeout / 1000)
+            local_state = self.get_local_state() or {}
+            self.set_local_state(local_state)
+
+    def stop(self) -> None:
+        """
+        Stops the awareness heartbeat task if it is running.
+        """
+        self._heartbeat_task.cancel()

packages/jupyter-ai/package.json

@@ -62,7 +62,7 @@
     "@emotion/react": "^11.10.5",
     "@emotion/styled": "^11.10.5",
     "@jupyter-notebook/application": "^7.2.0",
-    "@jupyter/chat": "^0.10.0",
+    "@jupyter/chat": "^0.11.0",
     "@jupyterlab/application": "^4.2.0",
     "@jupyterlab/apputils": "^4.2.0",
     "@jupyterlab/codeeditor": "^4.2.0",

packages/jupyter-ai/pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     # traitlets>=5.6 is required in JL4
     "traitlets>=5.6",
     "deepmerge>=2.0,<3",
-    "jupyterlab-chat>=0.10.0,<0.11.0",
+    "jupyterlab-chat>=0.11.0,<0.12.0",
 ]
 
 dynamic = ["version", "description", "authors", "urls", "keywords"]

yarn.lock

@@ -2237,7 +2237,7 @@ __metadata:
     "@emotion/react": ^11.10.5
     "@emotion/styled": ^11.10.5
     "@jupyter-notebook/application": ^7.2.0
-    "@jupyter/chat": ^0.10.0
+    "@jupyter/chat": ^0.11.0
     "@jupyterlab/application": ^4.2.0
     "@jupyterlab/apputils": ^4.2.0
     "@jupyterlab/builder": ^4.2.0
@@ -2322,9 +2322,9 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@jupyter/chat@npm:^0.10.0":
-  version: 0.10.0
-  resolution: "@jupyter/chat@npm:0.10.0"
+"@jupyter/chat@npm:^0.11.0":
+  version: 0.11.0
+  resolution: "@jupyter/chat@npm:0.11.0"
   dependencies:
     "@emotion/react": ^11.10.5
     "@emotion/styled": ^11.10.5
@@ -2346,7 +2346,7 @@ __metadata:
     clsx: ^2.1.0
     react: ^18.2.0
     react-dom: ^18.2.0
-  checksum: c9b0bc1c39b966a94415f6b95c86bda1a1bedefc654cec051ca6ebd30fc218a31cc1a62628bf2c5b529284cdafd5ce3fe9e15bb9bfdaa21960d47acdc2b66ebd
+  checksum: e7b5c47c94b6afa3ce6926c4f9ed1a913cfb95d7187015d058550b69cf26bdbd42240b8f550646f9ddfeed34c35ac81c7a664c1873152fa4ddd2a5f36ed5eb51
   languageName: node
   linkType: hard