Manages connection lifecycle, message polling, LLM orchestration, and persistence for a single user-agent session.
| 237 | |
| 238 | |
| 239 | class WebSocketChatHandler: |
| 240 | """Manages connection lifecycle, message polling, LLM orchestration, and persistence for a single user-agent session.""" |
| 241 | |
| 242 | def __init__( |
| 243 | self, |
| 244 | websocket: WebSocket, |
| 245 | agent_id: uuid.UUID, |
| 246 | token: str, |
| 247 | session_id: str | None = None, |
| 248 | lang: str = "en", |
| 249 | ): |
| 250 | self.websocket = websocket |
| 251 | self.agent_id = agent_id |
| 252 | self.token = token |
| 253 | self.session_id_param = session_id |
| 254 | self.lang = lang |
| 255 | |
| 256 | # State fields initialized during setup |
| 257 | self.user: User | None = None |
| 258 | self.agent: Agent | None = None |
| 259 | self.agent_name: str = "" |
| 260 | self.agent_type: str = "" |
| 261 | self.role_description: str = "" |
| 262 | self.welcome_message: str = "" |
| 263 | self.ctx_size: int = 100 |
| 264 | self.user_display_name: str = "" |
| 265 | self.llm_model: LLMModel | None = None |
| 266 | self.fallback_llm_model: LLMModel | None = None |
| 267 | self.conv_id: str | None = None |
| 268 | self.history_messages: list[ChatMessage] = [] |
| 269 | self.conversation: list[dict] = [] |
| 270 | self.current_user_text: str = "" |
| 271 | |
| 272 | async def run(self): |
| 273 | """Main entry point for handling the lifecycle of the WebSocket connection.""" |
| 274 | try: |
| 275 | # 1. Setup session (Authentication, permissions, loading models, history, etc.) |
| 276 | success = await self.setup() |
| 277 | if not success: |
| 278 | return |
| 279 | |
| 280 | # 2. Start the message receiving and processing loop |
| 281 | await self.message_loop() |
| 282 | |
| 283 | except WebSocketDisconnect: |
| 284 | logger.info(f"[WS] Client disconnected: {getattr(self.user, 'id', 'unknown')}") |
| 285 | await manager.disconnect(str(self.agent_id), self.websocket) |
| 286 | except Exception as e: |
| 287 | logger.exception(f"[WS] Unexpected error: {e}") |
| 288 | await manager.disconnect(str(self.agent_id), self.websocket) |
| 289 | |
| 290 | async def setup(self) -> bool: |
| 291 | """Accepts connection, authenticates user, verifies agent access, loads models, resolves session & history.""" |
| 292 | # Accept immediately so browser sees onopen without waiting for DB setup |
| 293 | await self.websocket.accept() |
| 294 | |
| 295 | # Authenticate |
| 296 | try: |