1+ """
2+ MessageRouter that manages message routing with callbacks.
3+
4+ This module provides a MessageRouter that:
5+ - Handles new chat connections
6+ - Routes slash commands and regular messages via callbacks
7+ - Manages lifecycle and cleanup
8+ """
9+
10+ from typing import Any , Callable , Dict , List , Set , TYPE_CHECKING
11+ from functools import partial
12+ from jupyterlab_chat .models import Message
13+ from pycrdt import ArrayEvent
14+ from traitlets .config import LoggingConfigurable
15+
16+ if TYPE_CHECKING :
17+ from jupyterlab_chat .ychat import YChat
18+
19+ from .utils import get_first_word
20+
21+
22+
23+ class MessageRouter (LoggingConfigurable ):
24+ """
25+ Router that manages ychat message routing.
26+
27+ The Router provides three callback points:
28+ 1. When new chats are initialized
29+ 2. When slash commands are received
30+ 3. When regular (non-slash) messages are received
31+ """
32+
33+ def __init__ (self , * args , ** kwargs ):
34+ super ().__init__ (* args , ** kwargs )
35+
36+ # Callback lists
37+ self .chat_init_observers : List [Callable [[str , "YChat" ], Any ]] = []
38+ self .slash_cmd_observers : Dict [str , List [Callable [[str , Message ], Any ]]] = {}
39+ self .chat_msg_observers : Dict [str , List [Callable [[str , Message ], Any ]]] = {}
40+
41+ # Active chat rooms
42+ self .active_chats : Dict [str , "YChat" ] = {}
43+
44+ # Root observers for keeping track of incoming messages
45+ self .message_observers : Dict [str , Callable ] = {}
46+
47+ def observe_chat_init (self , callback : Callable [[str , "YChat" ], Any ]) -> None :
48+ """
49+ Register a callback for when new chats are initialized.
50+
51+ Args:
52+ callback: Function called with (room_id: str, ychat: YChat) when chat connects
53+ """
54+ self .chat_init_observers .append (callback )
55+ self .log .info ("Registered new chat initialization callback" )
56+
57+ def observe_slash_cmd_msg (self , room_id : str , callback : Callable [[str , Message ], Any ]) -> None :
58+ """
59+ Register a callback for when slash commands are received.
60+
61+ Args:
62+ callback: Function called with (room_id: str, message: Message) for slash commands
63+ """
64+ if room_id not in self .slash_cmd_observers :
65+ self .slash_cmd_observers [room_id ] = []
66+
67+ self .slash_cmd_observers [room_id ].append (callback )
68+ self .log .info ("Registered slash command callback" )
69+
70+ def observe_chat_msg (self , room_id : str , callback : Callable [[str , Message ], Any ]) -> None :
71+ """
72+ Register a callback for when regular (non-slash) messages are received.
73+
74+ Args:
75+ callback: Function called with (room_id: str, message: Message) for regular messages
76+ """
77+ if room_id not in self .chat_msg_observers :
78+ self .chat_msg_observers [room_id ] = []
79+
80+ self .chat_msg_observers [room_id ].append (callback )
81+ self .log .info ("Registered message callback" )
82+
83+ def connect_chat (self , room_id : str , ychat : "YChat" ) -> None :
84+ """
85+ Connect a new chat session to the router.
86+
87+ Args:
88+ room_id: Unique identifier for the chat room
89+ ychat: YChat instance for the room
90+ """
91+ if room_id in self .active_chats :
92+ self .log .warning (f"Chat { room_id } )
93+ return
94+
95+ self .active_chats [room_id ] = ychat
96+
97+ # Set up message observer
98+ callback = partial (self ._on_message_change , room_id , ychat )
99+ ychat .ymessages .observe (callback )
100+ self .message_observers [room_id ] = callback
101+
102+ self .log .info (f"Connected chat { room_id } )
103+
104+ # Notify new chat observers
105+ self ._notify_chat_init_observers (room_id , ychat )
106+
107+ def disconnect_chat (self , room_id : str ) -> None :
108+ """
109+ Disconnect a chat session from the router.
110+
111+ Args:
112+ room_id: Unique identifier for the chat room
113+ """
114+ if room_id not in self .active_chats :
115+ return
116+
117+ # Remove message observer
118+ if room_id in self .message_observers :
119+ ychat = self .active_chats [room_id ]
120+ try :
121+ ychat .ymessages .unobserve (self .message_observers [room_id ])
122+ except Exception as e :
123+ self .log .warning (f"Failed to unobserve chat { room_id } { e } )
124+ del self .message_observers [room_id ]
125+
126+ del self .active_chats [room_id ]
127+ self .log .info (f"Disconnected chat { room_id } )
128+
129+ def _on_message_change (self , room_id : str , ychat : "YChat" , events : ArrayEvent ) -> None :
130+ """Handle incoming messages from YChat."""
131+ for change in events .delta : # type: ignore[attr-defined]
132+ if "insert" not in change .keys ():
133+ continue
134+
135+ # Process new messages (filter out raw_time duplicates)
136+ new_messages = [
137+ Message (** m ) for m in change ["insert" ]
138+ if not m .get ("raw_time" , False )
139+ ]
140+
141+ for message in new_messages :
142+ self ._route_message (room_id , message )
143+
144+ def _route_message (self , room_id : str , message : Message ) -> None :
145+ """
146+ Route an incoming message to appropriate observers.
147+
148+ Args:
149+ room_id: The chat room ID
150+ message: The message to route
151+ """
152+ first_word = get_first_word (message .body )
153+
154+ # Check if it's a slash command
155+ if first_word and first_word .startswith ("/" ):
156+ self ._notify_slash_cmd_observers (room_id , message )
157+ else :
158+ self ._notify_msg_observers (room_id , message )
159+
160+ def _notify_chat_init_observers (self , room_id : str , ychat : "YChat" ) -> None :
161+ """Notify all new chat observers."""
162+ for callback in self .chat_init_observers :
163+ try :
164+ callback (room_id , ychat )
165+ except Exception as e :
166+ self .log .error (f"New chat observer error for { room_id } { e } )
167+
168+ def _notify_slash_cmd_observers (self , room_id : str , message : Message ) -> None :
169+ """Notify all slash command observers."""
170+ callbacks = self .slash_cmd_observers .get (room_id , [])
171+ for callback in callbacks :
172+ try :
173+ callback (room_id , message )
174+ except Exception as e :
175+ self .log .error (f"Slash command observer error for { room_id } { e } )
176+
177+ def _notify_msg_observers (self , room_id : str , message : Message ) -> None :
178+ """Notify all message observers."""
179+ callbacks = self .chat_msg_observers .get (room_id , [])
180+ for callback in callbacks :
181+ try :
182+ callback (room_id , message )
183+ except Exception as e :
184+ self .log .error (f"Message observer error for { room_id } { e } )
185+
186+ def cleanup (self ) -> None :
187+ """Clean up router resources."""
188+ self .log .info ("Cleaning up MessageRouter..." )
189+
190+ # Disconnect all chats
191+ room_ids = list (self .active_chats .keys ())
192+ for room_id in room_ids :
193+ self .disconnect_chat (room_id )
194+
195+ # Clear callbacks
196+ self .chat_init_observers .clear ()
197+ self .slash_cmd_observers .clear ()
198+ self .chat_msg_observers .clear ()
199+
200+ self .log .info ("MessageRouter cleanup complete" )
0 commit comments