call dedicated method on Router when YChat is reset

Author: dlqqq

jupyter_ai_router/extension.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from typing import TYPE_CHECKING
 import time
 from jupyter_events import EventLogger
 from jupyter_server.extension.application import ExtensionApp
@@ -15,13 +16,18 @@
         "https://events.jupyter.org/jupyter_collaboration"
     )
 
+# Define `JSD_PRESENT` to indicate whether `jupyter_server_documents` is
+# installed in the current environment.
 JSD_PRESENT = False
 try:
     import jupyter_server_documents
     JSD_PRESENT = True
 except ImportError:
     pass
 
+if TYPE_CHECKING:
+    from jupyterlab_chat.ychat import YChat
+
 
 class RouterExtension(ExtensionApp):
     """
@@ -33,6 +39,8 @@ class RouterExtension(ExtensionApp):
         (r"jupyter-ai-router/health/?", RouteHandler),
     ]
 
+    router: MessageRouter
+
     def initialize_settings(self):
         """Initialize router settings and event listeners."""
         start = time.time()
@@ -79,8 +87,24 @@ async def _on_chat_event(
         # Connect chat to router
         self.router.connect_chat(room_id, ychat)
 
-    async def _get_chat(self, room_id: str):
-        """Get YChat instance for a room ID."""
+    async def _get_chat(self, room_id: str) -> YChat | None:
+        """
+        Get YChat instance for a room ID.
+
+        Dispatches to either `_get_chat_jcollab()` or `_get_chat_jsd()` based on
+        whether `jupyter_server_documents` is installed.
+        """
+
+        if JSD_PRESENT:
+            return await self._get_chat_jsd(room_id)
+        else:
+            return await self._get_chat_jcollab(room_id)
+    
+    async def _get_chat_jcollab(self, room_id: str) -> YChat | None:
+        """
+        Method used to retrieve the `YChat` instance for a given room when
+        `jupyter_server_documents` **is not** installed.
+        """
         if not self.serverapp:
             return None
 
@@ -91,6 +115,33 @@ async def _get_chat(self, room_id: str):
         except Exception as e:
             self.log.error(f"Error getting chat document for {room_id}: {e}")
             return None
+    
+    async def _get_chat_jsd(self, room_id: str) -> YChat | None:
+        """
+        Method used to retrieve the `YChat` instance for a given room when
+        `jupyter_server_documents` **is** installed.
+
+        This method uniquely attaches a callback which is fired whenever the
+        `YChat` is reset.
+        """
+        if not self.serverapp:
+            return None
+
+        try:
+            jcollab_api = self.serverapp.web_app.settings["jupyter_server_ydoc"]
+            yroom_manager = jcollab_api.yroom_manager
+            yroom = yroom_manager.get_room(room_id)
+            
+            def _on_ychat_reset(new_ychat: YChat):
+                self.router._on_chat_reset(room_id, new_ychat)
+
+            ychat = await yroom.get_jupyter_ydoc(on_reset=_on_ychat_reset)
+            return ychat
+        except Exception as e:
+            self.log.error(f"Error getting chat document for {room_id}: {e}")
+            return None
+
+
 
     async def stop_extension(self):
         """Clean up router when extension stops."""

jupyter_ai_router/router.py

@@ -55,6 +55,7 @@ def __init__(self, *args, **kwargs):
         self.chat_init_observers: List[Callable[[str, "YChat"], Any]] = []
         self.slash_cmd_observers: Dict[str, Dict[str, List[Callable[[str, str, Message], Any]]]] = {}
         self.chat_msg_observers: Dict[str, List[Callable[[str, Message], Any]]] = {}
+        self.chat_reset_observers: List[Callable[[str, "YChat"], Any]] = []
 
         # Active chat rooms
         self.active_chats: Dict[str, "YChat"] = {}
@@ -71,7 +72,7 @@ def observe_chat_init(self, callback: Callable[[str, "YChat"], Any]) -> None:
         """
         self.chat_init_observers.append(callback)
         self.log.info("Registered new chat initialization callback")
-
+    
     def observe_slash_cmd_msg(
         self, room_id: str, command_pattern: str, callback: Callable[[str, str, Message], Any]
     ) -> None:
@@ -230,6 +231,17 @@ def _notify_msg_observers(self, room_id: str, message: Message) -> None:
                 callback(room_id, message)
             except Exception as e:
                 self.log.error(f"Message observer error for {room_id}: {e}")
+    
+    def _on_chat_reset(self, room_id, ychat: "YChat") -> None:
+        """
+        Method to call when the YChat undergoes a document reset, e.g. when the
+        `.chat` file is modified directly on disk.
+        
+        NOTE: Document resets will only occur when `jupyter_server_documents` is
+        installed.
+        """
+        self.log.warning(f"Detected `YChat` document reset in room '{room_id}'.")
+        pass
 
     def cleanup(self) -> None:
         """Clean up router resources."""