Spaces:
Running
Running
| import asyncio | |
| import json | |
| import logging | |
| from collections.abc import Callable | |
| from urllib.parse import urlparse | |
| import aiohttp | |
| import websockets | |
| logger = logging.getLogger(__name__) | |
| class RoboticsClientCore: | |
| """Base client for LeRobot Arena robotics API""" | |
| def __init__(self, base_url: str = "http://localhost:8000"): | |
| self.base_url = base_url.rstrip("/") | |
| self.api_base = f"{self.base_url}/robotics" | |
| # WebSocket connection | |
| self.websocket: websockets.WebSocketServerProtocol | None = None | |
| self.workspace_id: str | None = None | |
| self.room_id: str | None = None | |
| self.role: str | None = None | |
| self.participant_id: str | None = None | |
| self.connected = False | |
| # Background task for message handling | |
| self._message_task: asyncio.Task | None = None | |
| # ============= REST API METHODS ============= | |
| async def list_rooms(self, workspace_id: str) -> list[dict]: | |
| """List all available rooms in a workspace""" | |
| async with aiohttp.ClientSession() as session: | |
| async with session.get( | |
| f"{self.api_base}/workspaces/{workspace_id}/rooms" | |
| ) as response: | |
| response.raise_for_status() | |
| result = await response.json() | |
| # Extract the rooms list from the response | |
| return result.get("rooms", []) | |
| async def create_room( | |
| self, workspace_id: str | None = None, room_id: str | None = None | |
| ) -> tuple[str, str]: | |
| """Create a new room and return (workspace_id, room_id)""" | |
| # Generate workspace ID if not provided | |
| final_workspace_id = workspace_id or self._generate_workspace_id() | |
| payload = {} | |
| if room_id: | |
| payload["room_id"] = room_id | |
| async with aiohttp.ClientSession() as session: | |
| async with session.post( | |
| f"{self.api_base}/workspaces/{final_workspace_id}/rooms", json=payload | |
| ) as response: | |
| response.raise_for_status() | |
| result = await response.json() | |
| return result["workspace_id"], result["room_id"] | |
| async def delete_room(self, workspace_id: str, room_id: str) -> bool: | |
| """Delete a room""" | |
| async with aiohttp.ClientSession() as session: | |
| async with session.delete( | |
| f"{self.api_base}/workspaces/{workspace_id}/rooms/{room_id}" | |
| ) as response: | |
| if response.status == 404: | |
| return False | |
| response.raise_for_status() | |
| result = await response.json() | |
| return result["success"] | |
| async def get_room_state(self, workspace_id: str, room_id: str) -> dict: | |
| """Get current room state""" | |
| async with aiohttp.ClientSession() as session: | |
| async with session.get( | |
| f"{self.api_base}/workspaces/{workspace_id}/rooms/{room_id}/state" | |
| ) as response: | |
| response.raise_for_status() | |
| result = await response.json() | |
| # Extract the state from the response | |
| return result.get("state", {}) | |
| async def get_room_info(self, workspace_id: str, room_id: str) -> dict: | |
| """Get basic room information""" | |
| async with aiohttp.ClientSession() as session: | |
| async with session.get( | |
| f"{self.api_base}/workspaces/{workspace_id}/rooms/{room_id}" | |
| ) as response: | |
| response.raise_for_status() | |
| result = await response.json() | |
| # Extract the room data from the response | |
| return result.get("room", {}) | |
| # ============= WEBSOCKET CONNECTION ============= | |
| async def connect_to_room( | |
| self, | |
| workspace_id: str, | |
| room_id: str, | |
| role: str, | |
| participant_id: str | None = None, | |
| ) -> bool: | |
| """Connect to a room as producer or consumer""" | |
| if self.connected: | |
| await self.disconnect() | |
| self.workspace_id = workspace_id | |
| self.room_id = room_id | |
| self.role = role | |
| self.participant_id = participant_id or f"{role}_{id(self)}" | |
| # Convert HTTP URL to WebSocket URL | |
| parsed = urlparse(self.base_url) | |
| ws_scheme = "wss" if parsed.scheme == "https" else "ws" | |
| ws_url = f"{ws_scheme}://{parsed.netloc}/robotics/workspaces/{workspace_id}/rooms/{room_id}/ws" | |
| initial_state_sync = None | |
| try: | |
| self.websocket = await websockets.connect(ws_url) | |
| # Send join message | |
| join_message = {"participant_id": self.participant_id, "role": role} | |
| await self.websocket.send(json.dumps(join_message)) | |
| # Wait for server response to join message | |
| try: | |
| response_text = await asyncio.wait_for( | |
| self.websocket.recv(), timeout=5.0 | |
| ) | |
| response = json.loads(response_text) | |
| if response.get("type") == "error": | |
| logger.error( | |
| f"Server rejected connection: {response.get('message')}" | |
| ) | |
| await self.websocket.close() | |
| return False | |
| if response.get("type") == "state_sync": | |
| # Consumer receives initial state sync, store it and wait for joined message | |
| logger.debug("Received initial state sync") | |
| initial_state_sync = response | |
| # Wait for the joined message | |
| response_text = await asyncio.wait_for( | |
| self.websocket.recv(), timeout=5.0 | |
| ) | |
| response = json.loads(response_text) | |
| if response.get("type") == "joined": | |
| logger.info(f"Successfully joined room {room_id} as {role}") | |
| elif response.get("type") == "error": | |
| logger.error( | |
| f"Server rejected connection: {response.get('message')}" | |
| ) | |
| await self.websocket.close() | |
| return False | |
| else: | |
| logger.warning(f"Unexpected response from server: {response}") | |
| elif response.get("type") == "joined": | |
| logger.info(f"Successfully joined room {room_id} as {role}") | |
| # Connection successful, continue with setup | |
| else: | |
| logger.warning(f"Unexpected response from server: {response}") | |
| except TimeoutError: | |
| logger.error("Timeout waiting for server response") | |
| await self.websocket.close() | |
| return False | |
| except json.JSONDecodeError: | |
| logger.error("Invalid JSON response from server") | |
| await self.websocket.close() | |
| return False | |
| # Start message handling task | |
| self._message_task = asyncio.create_task(self._handle_messages()) | |
| self.connected = True | |
| logger.info(f"Connected to room {room_id} as {role}") | |
| await self._on_connected() | |
| # Process initial state sync if we received one | |
| if initial_state_sync: | |
| await self._process_message(initial_state_sync) | |
| return True | |
| except Exception as e: | |
| logger.error(f"Failed to connect to room {room_id}: {e}") | |
| return False | |
| async def disconnect(self): | |
| """Disconnect from current room""" | |
| if self._message_task: | |
| self._message_task.cancel() | |
| try: | |
| await self._message_task | |
| except asyncio.CancelledError: | |
| pass | |
| self._message_task = None | |
| if self.websocket: | |
| await self.websocket.close() | |
| self.websocket = None | |
| self.connected = False | |
| self.workspace_id = None | |
| self.room_id = None | |
| self.role = None | |
| self.participant_id = None | |
| await self._on_disconnected() | |
| logger.info("Disconnected from room") | |
| # ============= MESSAGE HANDLING ============= | |
| async def _handle_messages(self): | |
| """Handle incoming WebSocket messages""" | |
| try: | |
| async for message in self.websocket: | |
| try: | |
| data = json.loads(message) | |
| await self._process_message(data) | |
| except json.JSONDecodeError: | |
| logger.error(f"Invalid JSON received: {message}") | |
| except Exception as e: | |
| logger.error(f"Error processing message: {e}") | |
| except websockets.exceptions.ConnectionClosed: | |
| logger.info("WebSocket connection closed") | |
| except Exception as e: | |
| logger.error(f"WebSocket error: {e}") | |
| finally: | |
| self.connected = False | |
| await self._on_disconnected() | |
| async def _process_message(self, data: dict): | |
| """Process incoming message based on type - to be overridden by subclasses""" | |
| msg_type = data.get("type") | |
| if msg_type == "joined": | |
| logger.info( | |
| f"Successfully joined room {data.get('room_id')} as {data.get('role')}" | |
| ) | |
| elif msg_type == "heartbeat_ack": | |
| logger.debug("Heartbeat acknowledged") | |
| else: | |
| # Let subclasses handle specific message types | |
| await self._handle_role_specific_message(data) | |
| async def _handle_role_specific_message(self, data: dict): | |
| """Handle role-specific messages - to be overridden by subclasses""" | |
| # ============= UTILITY METHODS ============= | |
| async def send_heartbeat(self): | |
| """Send heartbeat to server""" | |
| if not self.connected: | |
| return | |
| message = {"type": "heartbeat"} | |
| await self.websocket.send(json.dumps(message)) | |
| def is_connected(self) -> bool: | |
| """Check if client is connected""" | |
| return self.connected | |
| def get_connection_info(self) -> dict: | |
| """Get current connection information""" | |
| return { | |
| "connected": self.connected, | |
| "workspace_id": self.workspace_id, | |
| "room_id": self.room_id, | |
| "role": self.role, | |
| "participant_id": self.participant_id, | |
| "base_url": self.base_url, | |
| } | |
| # ============= HOOKS FOR SUBCLASSES ============= | |
| async def _on_connected(self): | |
| """Called when connection is established - to be overridden by subclasses""" | |
| async def _on_disconnected(self): | |
| """Called when connection is lost - to be overridden by subclasses""" | |
| # ============= CONTEXT MANAGER SUPPORT ============= | |
| async def __aenter__(self): | |
| return self | |
| async def __aexit__(self, exc_type, exc_val, exc_tb): | |
| await self.disconnect() | |
| # ============= WORKSPACE HELPERS ============= | |
| def _generate_workspace_id(self) -> str: | |
| """Generate a UUID-like workspace ID""" | |
| import uuid | |
| return str(uuid.uuid4()) | |
| class RoboticsProducer(RoboticsClientCore): | |
| """Producer client for controlling robots""" | |
| def __init__(self, base_url: str = "http://localhost:8000"): | |
| super().__init__(base_url) | |
| self._on_error_callback: Callable[[str], None] | None = None | |
| self._on_connected_callback: Callable[[], None] | None = None | |
| self._on_disconnected_callback: Callable[[], None] | None = None | |
| async def connect( | |
| self, workspace_id: str, room_id: str, participant_id: str | None = None | |
| ) -> bool: | |
| """Connect as producer to a room""" | |
| return await self.connect_to_room( | |
| workspace_id, room_id, "producer", participant_id | |
| ) | |
| # ============= PRODUCER METHODS ============= | |
| async def send_joint_update(self, joints: list[dict]): | |
| """Send joint updates""" | |
| if not self.connected: | |
| raise ValueError("Must be connected to send joint updates") | |
| message = {"type": "joint_update", "data": joints} | |
| await self.websocket.send(json.dumps(message)) | |
| async def send_state_sync(self, state: dict): | |
| """Send state synchronization (convert dict to list format)""" | |
| joints = [{"name": name, "value": value} for name, value in state.items()] | |
| await self.send_joint_update(joints) | |
| async def send_emergency_stop(self, reason: str = "Emergency stop"): | |
| """Send emergency stop signal""" | |
| if not self.connected: | |
| raise ValueError("Must be connected to send emergency stop") | |
| message = {"type": "emergency_stop", "reason": reason} | |
| await self.websocket.send(json.dumps(message)) | |
| # ============= EVENT CALLBACKS ============= | |
| def on_error(self, callback: Callable[[str], None]): | |
| """Set callback for error events""" | |
| self._on_error_callback = callback | |
| def on_connected(self, callback: Callable[[], None]): | |
| """Set callback for connection events""" | |
| self._on_connected_callback = callback | |
| def on_disconnected(self, callback: Callable[[], None]): | |
| """Set callback for disconnection events""" | |
| self._on_disconnected_callback = callback | |
| # ============= OVERRIDDEN HOOKS ============= | |
| async def _on_connected(self): | |
| if self._on_connected_callback: | |
| self._on_connected_callback() | |
| async def _on_disconnected(self): | |
| if self._on_disconnected_callback: | |
| self._on_disconnected_callback() | |
| async def _handle_role_specific_message(self, data: dict): | |
| """Handle producer-specific messages""" | |
| msg_type = data.get("type") | |
| if msg_type == "emergency_stop": | |
| logger.warning(f"🚨 Emergency stop: {data.get('reason', 'Unknown reason')}") | |
| if self._on_error_callback: | |
| self._on_error_callback( | |
| f"Emergency stop: {data.get('reason', 'Unknown reason')}" | |
| ) | |
| elif msg_type == "error": | |
| error_msg = data.get("message", "Unknown error") | |
| logger.error(f"Server error: {error_msg}") | |
| if self._on_error_callback: | |
| self._on_error_callback(error_msg) | |
| else: | |
| logger.warning(f"Unknown message type for producer: {msg_type}") | |
| class RoboticsConsumer(RoboticsClientCore): | |
| """Consumer client for receiving robot commands""" | |
| def __init__(self, base_url: str = "http://localhost:8000"): | |
| super().__init__(base_url) | |
| self._on_state_sync_callback: Callable[[dict], None] | None = None | |
| self._on_joint_update_callback: Callable[[list], None] | None = None | |
| self._on_error_callback: Callable[[str], None] | None = None | |
| self._on_connected_callback: Callable[[], None] | None = None | |
| self._on_disconnected_callback: Callable[[], None] | None = None | |
| async def connect( | |
| self, workspace_id: str, room_id: str, participant_id: str | None = None | |
| ) -> bool: | |
| """Connect as consumer to a room""" | |
| return await self.connect_to_room( | |
| workspace_id, room_id, "consumer", participant_id | |
| ) | |
| # ============= CONSUMER METHODS ============= | |
| async def get_state_sync(self) -> dict: | |
| """Get current state synchronously""" | |
| if not self.workspace_id or not self.room_id: | |
| raise ValueError("Must be connected to a room") | |
| state = await self.get_room_state(self.workspace_id, self.room_id) | |
| return state.get("joints", {}) | |
| # ============= EVENT CALLBACKS ============= | |
| def on_state_sync(self, callback: Callable[[dict], None]): | |
| """Set callback for state synchronization events""" | |
| self._on_state_sync_callback = callback | |
| def on_joint_update(self, callback: Callable[[list], None]): | |
| """Set callback for joint update events""" | |
| self._on_joint_update_callback = callback | |
| def on_error(self, callback: Callable[[str], None]): | |
| """Set callback for error events""" | |
| self._on_error_callback = callback | |
| def on_connected(self, callback: Callable[[], None]): | |
| """Set callback for connection events""" | |
| self._on_connected_callback = callback | |
| def on_disconnected(self, callback: Callable[[], None]): | |
| """Set callback for disconnection events""" | |
| self._on_disconnected_callback = callback | |
| # ============= OVERRIDDEN HOOKS ============= | |
| async def _on_connected(self): | |
| if self._on_connected_callback: | |
| self._on_connected_callback() | |
| async def _on_disconnected(self): | |
| if self._on_disconnected_callback: | |
| self._on_disconnected_callback() | |
| async def _handle_role_specific_message(self, data: dict): | |
| """Handle consumer-specific messages""" | |
| msg_type = data.get("type") | |
| if msg_type == "state_sync": | |
| if self._on_state_sync_callback: | |
| self._on_state_sync_callback(data.get("data", {})) | |
| elif msg_type == "joint_update": | |
| if self._on_joint_update_callback: | |
| self._on_joint_update_callback(data.get("data", [])) | |
| elif msg_type == "emergency_stop": | |
| logger.warning(f"🚨 Emergency stop: {data.get('reason', 'Unknown reason')}") | |
| if self._on_error_callback: | |
| self._on_error_callback( | |
| f"Emergency stop: {data.get('reason', 'Unknown reason')}" | |
| ) | |
| elif msg_type == "error": | |
| error_msg = data.get("message", "Unknown error") | |
| logger.error(f"Server error: {error_msg}") | |
| if self._on_error_callback: | |
| self._on_error_callback(error_msg) | |
| else: | |
| logger.warning(f"Unknown message type for consumer: {msg_type}") | |
| # ============= FACTORY FUNCTIONS ============= | |
| def create_client(role: str, base_url: str = "http://localhost:8000"): | |
| """Factory function to create the appropriate client based on role""" | |
| if role == "producer": | |
| return RoboticsProducer(base_url) | |
| if role == "consumer": | |
| return RoboticsConsumer(base_url) | |
| raise ValueError(f"Invalid role: {role}. Must be 'producer' or 'consumer'") | |
| async def create_producer_client( | |
| base_url: str = "http://localhost:8000", | |
| workspace_id: str | None = None, | |
| room_id: str | None = None, | |
| ) -> RoboticsProducer: | |
| """Create and connect a producer client""" | |
| client = RoboticsProducer(base_url) | |
| workspace_id, room_id = await client.create_room(workspace_id, room_id) | |
| await client.connect(workspace_id, room_id) | |
| return client | |
| async def create_consumer_client( | |
| workspace_id: str, room_id: str, base_url: str = "http://localhost:8000" | |
| ) -> RoboticsConsumer: | |
| """Create and connect a consumer client""" | |
| client = RoboticsConsumer(base_url) | |
| await client.connect(workspace_id, room_id) | |
| return client | |