Spaces:
Running
Running
File size: 3,190 Bytes
779968d 1ff744b 779968d a375433 779968d a375433 779968d | 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 | from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from optimization_manager import OptimizationManager
app = FastAPI()
PORT = 7860
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
print("starting")
await websocket.accept()
manager = OptimizationManager()
try:
while True:
# Receive message from client
message = await websocket.receive_json()
message_type = message.get("type")
# Handle different message types
if message_type == "init":
await handle_init(websocket, manager, message.get("settings", {}))
elif message_type == "reset":
await handle_reset(websocket, manager)
elif message_type == "next-step":
await handle_next_step(websocket, manager)
elif message_type == "prev-step":
await handle_prev_step(websocket, manager)
elif message_type == "play":
await handle_play(websocket, manager)
elif message_type == "pause":
await handle_pause(websocket, manager)
else:
await send_error(websocket, f"Unknown message type: {message_type}")
except WebSocketDisconnect:
print("Client disconnected")
except Exception as e:
print(f"Error: {e}")
await send_error(websocket, str(e))
async def handle_init(websocket: WebSocket, manager: OptimizationManager, settings: dict):
"""Handle settings update."""
data = manager.handle_update_settings(settings)
if data:
await websocket.send_json(data)
async def handle_reset(websocket: WebSocket, manager: OptimizationManager):
"""Handle trajectory reset."""
data = manager.handle_reset()
if data:
await websocket.send_json(data)
async def handle_next_step(websocket: WebSocket, manager: OptimizationManager):
"""Handle next step request."""
data = manager.handle_next_step()
if data:
await websocket.send_json(data)
async def handle_prev_step(websocket: WebSocket, manager: OptimizationManager):
"""Handle previous step request."""
data = manager.handle_prev_step()
if data:
await websocket.send_json(data)
async def handle_play(websocket: WebSocket, manager: OptimizationManager):
"""Handle play request (to be implemented)."""
# TODO: Implement play functionality with automatic step advancement
pass
async def handle_pause(websocket: WebSocket, manager: OptimizationManager):
"""Handle pause request (to be implemented)."""
# TODO: Implement pause functionality
pass
async def send_error(websocket: WebSocket, error_message: str):
"""Send error message to client."""
await websocket.send_json({
"type": "error",
"message": error_message
})
app.mount("/", StaticFiles(directory="../frontend/dist", html=True), name="frontend")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=PORT)
|