refactor: done phase 1

core/communication/models.py

@@ -0,0 +1,136 @@
+"""
+Message data structures for WebSocket communication.
+Based on worker.md protocol specification.
+"""
+from typing import Dict, Any, List, Optional, Union, Literal
+from pydantic import BaseModel, Field
+from datetime import datetime
+
+
+class SubscriptionObject(BaseModel):
+    """Individual camera subscription configuration."""
+    subscriptionIdentifier: str = Field(..., description="Format: displayId;cameraId")
+    rtspUrl: Optional[str] = Field(None, description="RTSP stream URL")
+    snapshotUrl: Optional[str] = Field(None, description="HTTP snapshot URL")
+    snapshotInterval: Optional[int] = Field(None, description="Snapshot interval in milliseconds")
+    modelUrl: str = Field(..., description="Pre-signed URL to .mpta file")
+    modelId: int = Field(..., description="Unique model identifier")
+    modelName: str = Field(..., description="Human-readable model name")
+    cropX1: Optional[int] = Field(None, description="Crop region X1 coordinate")
+    cropY1: Optional[int] = Field(None, description="Crop region Y1 coordinate")
+    cropX2: Optional[int] = Field(None, description="Crop region X2 coordinate")
+    cropY2: Optional[int] = Field(None, description="Crop region Y2 coordinate")
+
+
+class CameraConnection(BaseModel):
+    """Camera connection status for state reporting."""
+    subscriptionIdentifier: str
+    modelId: int
+    modelName: str
+    online: bool
+    cropX1: Optional[int] = None
+    cropY1: Optional[int] = None
+    cropX2: Optional[int] = None
+    cropY2: Optional[int] = None
+
+
+class DetectionData(BaseModel):
+    """Detection result data structure."""
+    detection: Dict[str, Any] = Field(..., description="Flat key-value detection results")
+    modelId: int
+    modelName: str
+
+
+# Incoming Messages from Backend to Worker
+
+class SetSubscriptionListMessage(BaseModel):
+    """Complete subscription list for declarative state management."""
+    type: Literal["setSubscriptionList"] = "setSubscriptionList"
+    subscriptions: List[SubscriptionObject]
+
+
+class SetSessionIdPayload(BaseModel):
+    """Session ID association payload."""
+    displayIdentifier: str
+    sessionId: Optional[int] = None
+
+
+class SetSessionIdMessage(BaseModel):
+    """Associate session ID with display."""
+    type: Literal["setSessionId"] = "setSessionId"
+    payload: SetSessionIdPayload
+
+
+class SetProgressionStagePayload(BaseModel):
+    """Progression stage payload."""
+    displayIdentifier: str
+    progressionStage: Optional[str] = None
+
+
+class SetProgressionStageMessage(BaseModel):
+    """Set progression stage for display."""
+    type: Literal["setProgressionStage"] = "setProgressionStage"
+    payload: SetProgressionStagePayload
+
+
+class RequestStateMessage(BaseModel):
+    """Request current worker state."""
+    type: Literal["requestState"] = "requestState"
+
+
+class PatchSessionResultPayload(BaseModel):
+    """Patch session result payload."""
+    sessionId: int
+    success: bool
+    message: str
+
+
+class PatchSessionResultMessage(BaseModel):
+    """Response to patch session request."""
+    type: Literal["patchSessionResult"] = "patchSessionResult"
+    payload: PatchSessionResultPayload
+
+
+# Outgoing Messages from Worker to Backend
+
+class StateReportMessage(BaseModel):
+    """Periodic heartbeat with system metrics."""
+    type: Literal["stateReport"] = "stateReport"
+    cpuUsage: float
+    memoryUsage: float
+    gpuUsage: Optional[float] = None
+    gpuMemoryUsage: Optional[float] = None
+    cameraConnections: List[CameraConnection]
+
+
+class ImageDetectionMessage(BaseModel):
+    """Detection event message."""
+    type: Literal["imageDetection"] = "imageDetection"
+    subscriptionIdentifier: str
+    timestamp: str = Field(default_factory=lambda: datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"))
+    sessionId: Optional[int] = None
+    data: DetectionData
+
+
+class PatchSessionMessage(BaseModel):
+    """Request to modify session data."""
+    type: Literal["patchSession"] = "patchSession"
+    sessionId: int
+    data: Dict[str, Any] = Field(..., description="Partial DisplayPersistentData structure")
+
+
+# Union type for all incoming messages
+IncomingMessage = Union[
+    SetSubscriptionListMessage,
+    SetSessionIdMessage,
+    SetProgressionStageMessage,
+    RequestStateMessage,
+    PatchSessionResultMessage
+]
+
+# Union type for all outgoing messages
+OutgoingMessage = Union[
+    StateReportMessage,
+    ImageDetectionMessage,
+    PatchSessionMessage
+]