Spaces:
Running
Running
| from fastapi import FastAPI, WebSocket, WebSocketDisconnect | |
| from fastapi.staticfiles import StaticFiles | |
| from optimization_manager import OptimizationManager | |
| app = FastAPI() | |
| PORT = 7860 | |
| 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) | |