AI实验室后端

bbit_ai/app/db/postgres/ws_manager.py

@@ -4,13 +4,17 @@ from uuid import UUID
 
 from fastapi import WebSocket
 
+PROJ_IOT_ONLINE_STATE = 0
+PROJ_SENTINEL_VEHICLE_STATES = 1
+PROJ_SENTINEL_MONITOR_STATUS = 2
+
 
 class ConnectionManager:
     def __init__(self):
         self.active_connections: List[dict] = []  # 保存 websocket 和用户信息
         self.lock = asyncio.Lock()
 
-    # proj_id:0：在线状态 1：畜牧车辆进入
+    # proj_id:0：在线状态 1：畜牧车辆进入 2：畜牧监控大屏
     async def connect(
         self, websocket: WebSocket, user_id: UUID, dept_id: str, proj_id: int
     ):
@@ -24,6 +28,9 @@ class ConnectionManager:
                     "proj_id": proj_id,
                 }
             )
+            print(
+                "Device" + str(user_id) + " in project " + str(proj_id) + " connected"
+            )
 
     async def disconnect(self, websocket: WebSocket):
         async with self.lock:
@@ -34,7 +41,7 @@ class ConnectionManager:
     async def noticeOnlineStatus(self, message: dict):
         async with self.lock:
             for conn in self.active_connections:
-                if conn["proj_id"] == 0:
+                if conn["proj_id"] == PROJ_IOT_ONLINE_STATE:
                     await conn["ws"].send_json(message)
 
     async def noticeSentinel(
@@ -46,5 +53,23 @@ class ConnectionManager:
         async with self.lock:
             for conn in self.active_connections:
                 if target_departments:
-                    if conn["proj_id"] == 1 and conn["dept_id"] in target_departments:
+                    if (
+                        conn["proj_id"] == PROJ_SENTINEL_VEHICLE_STATES
+                        and conn["dept_id"] in target_departments
+                    ):
+                        await conn["ws"].send_json(message)
+
+    async def noticeSentinelMonitorStatus(
+        self, message: dict, target_departments: List[UUID] = None
+    ):
+        """
+        target_departments: 指定哪些部门能收到消息
+        """
+        async with self.lock:
+            for conn in self.active_connections:
+                if target_departments:
+                    if (
+                        conn["proj_id"] == PROJ_SENTINEL_MONITOR_STATUS
+                        and conn["dept_id"] in target_departments
+                    ):
                         await conn["ws"].send_json(message)