完成训练模块的转移

core/websocket_app.py

@@ -0,0 +1,30 @@
+from fastapi import WebSocket
+from fastapi import FastAPI
+from starlette.websockets import WebSocketState
+
+from utils.websocket_server import room_manager
+
+
+def websocket_config(app: FastAPI):
+
+    @app.websocket("/{room}")
+    async def websocket_room(websocket: WebSocket, room: str):
+        """
+        websocket 房间管理
+        :param websocket:
+        :param room:
+        :return:
+        """
+        await websocket.accept()
+        await room_manager.add_to_room(room, websocket)
+        try:
+            while True:
+                data = await websocket.receive_text()
+                await room_manager.broadcast_to_room(room, data, exclude_websocket=websocket)
+        except Exception as e:
+            if websocket.client_state != WebSocketState.DISCONNECTED:
+                await websocket.close(code=1000)
+        finally:
+            await room_manager.remove_from_room(room, websocket)
+            if websocket.client_state != WebSocketState.DISCONNECTED:
+                await websocket.close(code=1001)