mirror of
https://github.com/OpenBMB/ChatDev.git
synced 2026-04-25 11:18:06 +00:00
82 lines
3.1 KiB
Python
Executable File
82 lines
3.1 KiB
Python
Executable File
import logging
|
|
from typing import Any, Dict
|
|
|
|
from utils.exceptions import ValidationError
|
|
|
|
from server.services.session_execution import SessionExecutionController
|
|
from server.services.session_store import WorkflowSessionStore
|
|
|
|
|
|
class MessageHandler:
|
|
"""Routes WebSocket messages to the appropriate handlers."""
|
|
|
|
def __init__(
|
|
self,
|
|
session_store: WorkflowSessionStore,
|
|
session_controller: SessionExecutionController,
|
|
workflow_run_service=None,
|
|
) -> None:
|
|
self.session_store = session_store
|
|
self.session_controller = session_controller
|
|
self.workflow_run_service = workflow_run_service
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
async def handle_message(self, session_id: str, data: Dict[str, Any], websocket_manager):
|
|
message_type = data.get("type")
|
|
if message_type == "human_input":
|
|
await self._handle_human_input(session_id, data, websocket_manager)
|
|
elif message_type == "ping":
|
|
await self._handle_ping(session_id, websocket_manager)
|
|
elif message_type == "get_status":
|
|
await self._handle_get_status(session_id, websocket_manager)
|
|
else:
|
|
await websocket_manager.send_message(
|
|
session_id,
|
|
{"type": "error", "data": {"message": f"Unknown message type: {message_type}"}},
|
|
)
|
|
|
|
async def _handle_human_input(self, session_id: str, data: Dict[str, Any], websocket_manager):
|
|
try:
|
|
payload = data.get("data", {}) or {}
|
|
user_input = payload.get("input", "")
|
|
attachments = payload.get("attachments") or []
|
|
|
|
if not user_input and not attachments:
|
|
await websocket_manager.send_message(
|
|
session_id,
|
|
{"type": "error", "data": {"message": "Empty input"}},
|
|
)
|
|
return
|
|
|
|
self.session_controller.provide_human_input(
|
|
session_id,
|
|
{"text": user_input, "attachments": attachments},
|
|
)
|
|
|
|
await websocket_manager.send_message(
|
|
session_id,
|
|
{"type": "input_received", "data": {"message": "Input received"}},
|
|
)
|
|
|
|
except ValidationError as exc:
|
|
await websocket_manager.send_message(
|
|
session_id,
|
|
{"type": "error", "data": {"message": str(exc)}},
|
|
)
|
|
except Exception as exc:
|
|
self.logger.error("Error handling human input for session %s: %s", session_id, exc)
|
|
await websocket_manager.send_message(
|
|
session_id,
|
|
{"type": "error", "data": {"message": str(exc)}},
|
|
)
|
|
|
|
async def _handle_ping(self, session_id: str, websocket_manager):
|
|
await websocket_manager.handle_heartbeat(session_id)
|
|
|
|
async def _handle_get_status(self, session_id: str, websocket_manager):
|
|
session_info = self.session_store.get_session_info(session_id)
|
|
await websocket_manager.send_message(
|
|
session_id,
|
|
{"type": "status", "data": session_info or {"message": "Session not found"}},
|
|
)
|