代码内容转移

backend_service/src/main.py

@@ -0,0 +1,93 @@
+import asyncio
+import os
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.staticfiles import StaticFiles
+import logging
+import threading
+import rclpy
+
+from .models import GeneratePlanRequest, ExecuteMissionRequest
+from .websocket_manager import websocket_manager
+from .py_tree_generator import py_tree_generator
+from .ros2_client import MissionActionClient
+
+# --- Application Setup ---
+app = FastAPI(
+    title="Drone Backend Service",
+    description="Handles mission planning, generation, and execution for the drone.",
+    version="1.0.0",
+)
+
+# --- Mount Static Files for Visualizations ---
+static_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'generated_visualizations'))
+app.mount("/static", StaticFiles(directory=static_dir), name="static")
+
+# --- ROS2 Node and Client Initialization ---
+rclpy.init()
+ros2_client = MissionActionClient()
+
+def run_ros2_node():
+    """Spins the ROS2 node in a dedicated thread."""
+    logging.info("Starting to spin ROS2 node...")
+    rclpy.spin(ros2_client)
+    logging.info("ROS2 node has stopped spinning.")
+
+# --- API Endpoints ---
+
+@app.post("/generate_plan", response_model=dict)
+async def generate_plan_endpoint(request: GeneratePlanRequest):
+    """
+    Receives a user prompt and returns a generated `py_tree.json` with a visualization URL.
+    """
+    try:
+        pytree_dict = await py_tree_generator.generate(request.user_prompt)
+        return pytree_dict
+    except RuntimeError as e:
+        return {"error": str(e)}
+
+@app.post("/execute_mission", response_model=dict)
+async def execute_mission_endpoint(request: ExecuteMissionRequest):
+    """
+    Receives a `py_tree.json` and sends it to the drone for execution.
+    """
+    ros2_client.send_goal(request.py_tree)
+    return {"status": "execution_started"}
+
+@app.websocket("/ws/status")
+async def websocket_endpoint(websocket: WebSocket):
+    """
+    Handles the WebSocket connection for real-time status updates.
+    """
+    await websocket_manager.connect(websocket)
+    try:
+        while True:
+            await websocket.receive_text()
+    except WebSocketDisconnect:
+        websocket_manager.disconnect(websocket)
+        logging.info("Client disconnected from WebSocket.")
+
+
+# --- Server Lifecycle ---
+
+@app.on_event("startup")
+async def startup_event():
+    """
+    On startup, get the current asyncio event loop and pass it to the websocket manager.
+    Also, start the ROS2 node in a background thread.
+    """
+    # Configure WebSocket Manager
+    loop = asyncio.get_running_loop()
+    websocket_manager.set_loop(loop)
+    logging.info("WebSocket event loop configured.")
+
+    # Start ROS2 node in a background thread
+    ros2_thread = threading.Thread(target=run_ros2_node, daemon=True)
+    ros2_thread.start()
+    logging.info("ROS2 node thread started.")
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    logging.info("Backend service shutting down.")
+    ros2_client.destroy_node()
+    rclpy.shutdown()
+    logging.info("ROS2 node shut down successfully.")