94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
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.")
|