Spaces:
Running
Running
File size: 5,030 Bytes
02eac4b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
import logging
from enum import Enum
from pydantic import BaseModel, Field
logger = logging.getLogger(__name__)
# ============= ENUMS =============
class ParticipantRole(str, Enum):
"""Participant roles in a robotics room"""
PRODUCER = "producer" # Controller/master
CONSUMER = "consumer" # Robot/visualizer
class MessageType(str, Enum):
"""WebSocket message types organized by category"""
# === CONNECTION & LIFECYCLE ===
JOINED = "joined" # Confirmation of successful room join
ERROR = "error" # Error notifications
# === CONNECTION HEALTH ===
HEARTBEAT = "heartbeat" # Client ping for connection health
HEARTBEAT_ACK = "heartbeat_ack" # Server response to heartbeat
# === ROBOT CONTROL (Core) ===
JOINT_UPDATE = "joint_update" # Producer → Consumers: Joint position commands
STATE_SYNC = "state_sync" # Server → Consumer: Initial state synchronization
# === EMERGENCY & SAFETY ===
EMERGENCY_STOP = "emergency_stop" # Emergency stop command (any direction)
# === STATUS & MONITORING ===
STATUS_UPDATE = "status_update" # General status updates
CONSUMER_STATUS = "consumer_status" # Specific consumer status reports
class RobotType(str, Enum):
"""Supported robot types"""
SO_ARM_100 = "so-arm100"
GENERIC = "generic"
# ============= CORE MODELS =============
class CreateRoomRequest(BaseModel):
"""Request to create a new robotics room"""
room_id: str | None = None
workspace_id: str | None = None # Optional - will be generated if not provided
name: str | None = None
robot_type: RobotType = RobotType.SO_ARM_100
max_consumers: int = Field(default=10, ge=1, le=100)
class JointUpdate(BaseModel):
"""Single joint position update"""
name: str = Field(..., min_length=1, max_length=50)
value: float = Field(..., ge=-360.0, le=360.0) # Degrees
speed: float | None = Field(None, ge=0.0, le=100.0) # Percentage
class JointCommand(BaseModel):
"""Command containing multiple joint updates"""
joints: list[JointUpdate] = Field(..., min_length=1, max_length=20)
class EmergencyStop(BaseModel):
"""Emergency stop command"""
enabled: bool = True
# ============= RESPONSE MODELS =============
class ParticipantInfo(BaseModel):
"""Information about room participants"""
producer: str | None
consumers: list[str]
total: int
class RoomInfo(BaseModel):
"""Basic room information"""
id: str
workspace_id: str
participants: ParticipantInfo
joints_count: int
has_producer: bool
active_consumers: int
class RoomState(BaseModel):
"""Detailed room state"""
room_id: str
workspace_id: str
joints: dict[str, float]
participants: ParticipantInfo
timestamp: str
# ============= WEBSOCKET MODELS =============
class WebSocketMessage(BaseModel):
"""Base WebSocket message"""
type: MessageType
timestamp: str | None = None
class JoinMessage(BaseModel):
"""Message to join a room"""
participant_id: str = Field(..., min_length=1, max_length=100)
role: ParticipantRole
# === CONNECTION & LIFECYCLE MESSAGES ===
class JoinedMessage(WebSocketMessage):
"""Confirmation of successful room join"""
type: MessageType = MessageType.JOINED
room_id: str
role: ParticipantRole
class ErrorMessage(WebSocketMessage):
"""Error notification message"""
type: MessageType = MessageType.ERROR
message: str
code: str | None = None
# === CONNECTION HEALTH MESSAGES ===
class HeartbeatMessage(WebSocketMessage):
"""Heartbeat ping from client"""
type: MessageType = MessageType.HEARTBEAT
class HeartbeatAckMessage(WebSocketMessage):
"""Heartbeat acknowledgment from server"""
type: MessageType = MessageType.HEARTBEAT_ACK
# === ROBOT CONTROL MESSAGES ===
class JointUpdateMessage(WebSocketMessage):
"""Joint position update command (Producer → Consumers)"""
type: MessageType = MessageType.JOINT_UPDATE
data: list[JointUpdate]
source: str | None = None # "api" or participant_id
class StateSyncMessage(WebSocketMessage):
"""Initial state synchronization (Server → Consumer)"""
type: MessageType = MessageType.STATE_SYNC
data: dict[str, float] # Current joint positions
# === EMERGENCY & SAFETY MESSAGES ===
class EmergencyStopMessage(WebSocketMessage):
"""Emergency stop command"""
type: MessageType = MessageType.EMERGENCY_STOP
enabled: bool = True
reason: str | None = None
# === STATUS & MONITORING MESSAGES ===
class StatusUpdateMessage(WebSocketMessage):
"""General status update"""
type: MessageType = MessageType.STATUS_UPDATE
status: str
data: dict | None = None
class ConsumerStatusMessage(WebSocketMessage):
"""Consumer-specific status report"""
type: MessageType = MessageType.CONSUMER_STATUS
consumer_id: str
status: str # "ready", "busy", "error", "offline"
joint_states: list[dict] | None = None
|