Core message routing layer for Jupyter AI

README.md

@@ -2,28 +2,46 @@
 
 [![Github Actions Status](https://github.com/jupyter-ai-contrib/jupyter-ai-router/workflows/Build/badge.svg)](https://github.com/jupyter-ai-contrib/jupyter-ai-router/actions/workflows/build.yml)
 
-Core routing layer of Jupyter AI
+Core message routing layer for Jupyter AI
 
-This extension is composed of a Python package named `jupyter_ai_router`
-for the server extension and a NPM package named `@jupyter-ai/router`
-for the frontend extension.
+This extension provides the foundational message routing functionality for Jupyter AI. It automatically detects new chat sessions and routes messages to registered callbacks based on message type (slash commands vs regular messages). Extensions can register callbacks to handle specific chat events without needing to manage chat lifecycle directly.
 
-## QUICK START
+## Usage
 
-Everything that follows after this section was from the extension template. We
-will need to revise the rest of this README.
+### Basic MessageRouter Setup
 
-Development install:
+```python
+# The router is available in other extensions via settings
+router = self.serverapp.web_app.settings.get("jupyter-ai", {}).get("router")
 
+# Register callbacks for different event types
+def on_new_chat(room_id: str, ychat: YChat):
+    print(f"New chat connected: {room_id}")
+
+def on_slash_command(room_id: str, message: Message):
+    print(f"Slash command in {room_id}: {message.body}")
+
+def on_regular_message(room_id: str, message: Message):`
+    print(f"Regular message in {room_id}: {message.body}")
+
+# Register the callbacks
+router.observe_chat_init(on_new_chat)
+router.observe_slash_cmd_msg("room-id", on_slash_command)
+router.observe_chat_msg("room-id", on_regular_message)
 ```
-micromamba install uv jupyterlab nodejs=22
-jlpm
-jlpm dev:install
-```
 
-## Requirements
+### Message Flow
+
+1. **Router detects new chats** - Automatically listens for chat room initialization events
+2. **Router connects chats** - Establishes observers on YChat message streams
+3. **Router routes messages** - Calls appropriate callbacks based on message type (slash vs regular)
+4. **Extensions respond** - Your callbacks receive room_id and message data
+
+### Available Methods
 
-- JupyterLab >= 4.0.0
+- `observe_chat_init(callback)` - Called when new chat sessions are initialized with `(room_id, ychat)`
+- `observe_slash_cmd_msg(room_id, callback)` - Called for messages starting with `/` in a specific room
+- `observe_chat_msg(room_id, callback)` - Called for regular (non-slash) messages in a specific room
 
 ## Install
 

jupyter_ai_router/__init__.py

@@ -7,7 +7,8 @@
     import warnings
     warnings.warn("Importing 'jupyter_ai_router' outside a proper installation.")
     __version__ = "dev"
-from .handlers import setup_handlers
+
+from .extension import RouterExtension
 
 
 def _jupyter_labextension_paths():
@@ -18,19 +19,4 @@ def _jupyter_labextension_paths():
 
 
 def _jupyter_server_extension_points():
-    return [{
-        "module": "jupyter_ai_router"
-    }]
-
-
-def _load_jupyter_server_extension(server_app):
-    """Registers the API handler to receive HTTP requests from the frontend extension.
-
-    Parameters
-    ----------
-    server_app: jupyterlab.labapp.LabApp
-        JupyterLab application instance
-    """
-    setup_handlers(server_app.web_app)
-    name = "jupyter_ai_router"
-    server_app.log.info(f"Registered {name} server extension")
+    return [{"module": "jupyter_ai_router", "app": RouterExtension}]

jupyter_ai_router/extension.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+import time
+from jupyter_events import EventLogger
+from jupyter_server.extension.application import ExtensionApp
+
+from jupyter_ai_router.handlers import RouteHandler
+
+from .router import MessageRouter
+
+# Check jupyter-collaboration version for compatibility
+try:
+    from jupyter_collaboration import __version__ as jupyter_collaboration_version
+
+    JCOLLAB_VERSION = int(jupyter_collaboration_version[0])
+    if JCOLLAB_VERSION >= 3:
+        from jupyter_server_ydoc.utils import JUPYTER_COLLABORATION_EVENTS_URI
+    else:
+        from jupyter_collaboration.utils import JUPYTER_COLLABORATION_EVENTS_URI
+except ImportError:
+    # Fallback if jupyter-collaboration is not available
+    JUPYTER_COLLABORATION_EVENTS_URI = (
+        "https://events.jupyter.org/jupyter_collaboration"
+    )
+
+
+class RouterExtension(ExtensionApp):
+    """
+    Jupyter AI Router Extension
+    """
+
+    name = "jupyter_ai_router"
+    handlers = [
+        (r"jupyter-ai-router/health/?", RouteHandler),
+    ]
+
+    def initialize_settings(self):
+        """Initialize router settings and event listeners."""
+        start = time.time()
+
+        # Create MessageRouter instance
+        self.router = MessageRouter(parent=self)
+
+        # Make router available to other extensions
+        if "jupyter-ai" not in self.settings:
+            self.settings["jupyter-ai"] = {}
+        self.settings["jupyter-ai"]["router"] = self.router
+
+        # Listen for new chat room events
+        if self.serverapp is not None:
+            self.event_logger = self.serverapp.web_app.settings["event_logger"]
+            self.event_logger.add_listener(
+                schema_id=JUPYTER_COLLABORATION_EVENTS_URI, listener=self._on_chat_event
+            )
+
+        elapsed = time.time() - start
+        self.log.info(f"Initialized RouterExtension in {elapsed:.2f}s")
+
+    async def _on_chat_event(
+        self, logger: EventLogger, schema_id: str, data: dict
+    ) -> None:
+        """Handle chat room events and connect new chats to router."""
+        # Only handle chat room initialization events
+        if not (
+            data["room"].startswith("text:chat:")
+            and data["action"] == "initialize"
+            and data["msg"] == "Room initialized"
+        ):
+            return
+
+        room_id = data["room"]
+        self.log.info(f"New chat room detected: {room_id}")
+
+        # Get YChat document for the room
+        ychat = await self._get_chat(room_id)
+        if ychat is None:
+            self.log.error(f"Failed to get YChat for room {room_id}")
+            return
+
+        # 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."""
+        if not self.serverapp:
+            return None
+
+        try:
+            if JCOLLAB_VERSION >= 3:
+                collaboration = self.serverapp.web_app.settings["jupyter_server_ydoc"]
+                document = await collaboration.get_document(room_id=room_id, copy=False)
+            else:
+                collaboration = self.serverapp.web_app.settings["jupyter_collaboration"]
+                server = collaboration.ywebsocket_server
+                room = await server.get_room(room_id)
+                document = room._document
+
+            return document
+        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."""
+        try:
+            if hasattr(self, "router"):
+                self.router.cleanup()
+        except Exception as e:
+            self.log.error(f"Error during router cleanup: {e}")

jupyter_ai_router/handlers.py

@@ -1,7 +1,6 @@
 import json
 
 from jupyter_server.base.handlers import APIHandler
-from jupyter_server.utils import url_path_join
 import tornado
 
 class RouteHandler(APIHandler):
@@ -11,14 +10,5 @@ class RouteHandler(APIHandler):
     @tornado.web.authenticated
     def get(self):
         self.finish(json.dumps({
-            "data": "This is /jupyter-ai-router/get-example endpoint!"
+            "data": "JupyterLab extension @jupyter-ai/router is activated!"
         }))
-
-
-def setup_handlers(web_app):
-    host_pattern = ".*$"
-
-    base_url = web_app.settings["base_url"]
-    route_pattern = url_path_join(base_url, "jupyter-ai-router", "get-example")
-    handlers = [(route_pattern, RouteHandler)]
-    web_app.add_handlers(host_pattern, handlers)