Improve api granularity to empower MCP server (#63)

Author: GabrielBarberini

.github/workflows/pylint.yml

@@ -21,4 +21,4 @@ jobs:
         pip install -r requirements.txt
     - name: Analysing the code with pylint
       run: |
-        pylint -d C0415,C0200,C0301,C0114,R0903,C0115,W0246,R0914,C0209,E1121,C0103,C2801,R0801,E1101,E0401,E0611,R0911,C0116,W0212,W0719,W0601,W1203,W0123,W0511,W0621,R0913,R0917 $(git ls-files '*.py')
+        pylint -d R0912,C0415,C0200,C0301,C0114,R0903,C0115,W0246,R0914,C0209,E1121,C0103,C2801,R0801,E1101,E0401,E0611,R0911,C0116,W0212,W0719,W0601,W1203,W0123,W0511,W0621,R0913,R0917 $(git ls-files '*.py')

src/controllers/flight.py

@@ -1,11 +1,17 @@
+from fastapi import HTTPException, status
+
 from src.controllers.interface import (
     ControllerBase,
     controller_exception_handler,
 )
-from src.views.flight import FlightSimulation
-from src.models.flight import FlightModel
+from src.views.flight import FlightSimulation, FlightCreated
+from src.models.flight import (
+    FlightModel,
+    FlightWithReferencesRequest,
+)
 from src.models.environment import EnvironmentModel
 from src.models.rocket import RocketModel
+from src.repositories.interface import RepositoryInterface
 from src.services.flight import FlightService
 
 
@@ -21,6 +27,56 @@ class FlightController(ControllerBase):
     def __init__(self):
         super().__init__(models=[FlightModel])
 
+    async def _load_environment(self, environment_id: str) -> EnvironmentModel:
+        repo_cls = RepositoryInterface.get_model_repo(EnvironmentModel)
+        async with repo_cls() as repo:
+            environment = await repo.read_environment_by_id(environment_id)
+        if environment is None:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail="Environment not found",
+            )
+        return environment
+
+    async def _load_rocket(self, rocket_id: str) -> RocketModel:
+        repo_cls = RepositoryInterface.get_model_repo(RocketModel)
+        async with repo_cls() as repo:
+            rocket = await repo.read_rocket_by_id(rocket_id)
+        if rocket is None:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail="Rocket not found",
+            )
+        return rocket
+
+    @controller_exception_handler
+    async def create_flight_from_references(
+        self, payload: FlightWithReferencesRequest
+    ) -> FlightCreated:
+        environment = await self._load_environment(payload.environment_id)
+        rocket = await self._load_rocket(payload.rocket_id)
+        flight_model = payload.flight.assemble(
+            environment=environment,
+            rocket=rocket,
+        )
+        return await self.post_flight(flight_model)
+
+    @controller_exception_handler
+    async def update_flight_from_references(
+        self,
+        flight_id: str,
+        payload: FlightWithReferencesRequest,
+    ) -> None:
+        environment = await self._load_environment(payload.environment_id)
+        rocket = await self._load_rocket(payload.rocket_id)
+        flight_model = payload.flight.assemble(
+            environment=environment,
+            rocket=rocket,
+        )
+        flight_model.set_id(flight_id)
+        await self.put_flight_by_id(flight_id, flight_model)
+        return
+
     @controller_exception_handler
     async def update_environment_by_flight_id(
         self, flight_id: str, *, environment: EnvironmentModel

src/controllers/rocket.py

@@ -1,9 +1,16 @@
+from fastapi import HTTPException, status
+
 from src.controllers.interface import (
     ControllerBase,
     controller_exception_handler,
 )
-from src.views.rocket import RocketSimulation
-from src.models.rocket import RocketModel
+from src.views.rocket import RocketSimulation, RocketCreated
+from src.models.motor import MotorModel
+from src.models.rocket import (
+    RocketModel,
+    RocketWithMotorReferenceRequest,
+)
+from src.repositories.interface import RepositoryInterface
 from src.services.rocket import RocketService
 
 
@@ -19,6 +26,37 @@ class RocketController(ControllerBase):
     def __init__(self):
         super().__init__(models=[RocketModel])
 
+    async def _load_motor(self, motor_id: str) -> MotorModel:
+        repo_cls = RepositoryInterface.get_model_repo(MotorModel)
+        async with repo_cls() as repo:
+            motor = await repo.read_motor_by_id(motor_id)
+        if motor is None:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail="Motor not found",
+            )
+        return motor
+
+    @controller_exception_handler
+    async def create_rocket_from_motor_reference(
+        self, payload: RocketWithMotorReferenceRequest
+    ) -> RocketCreated:
+        motor = await self._load_motor(payload.motor_id)
+        rocket_model = payload.rocket.assemble(motor)
+        return await self.post_rocket(rocket_model)
+
+    @controller_exception_handler
+    async def update_rocket_from_motor_reference(
+        self,
+        rocket_id: str,
+        payload: RocketWithMotorReferenceRequest,
+    ) -> None:
+        motor = await self._load_motor(payload.motor_id)
+        rocket_model = payload.rocket.assemble(motor)
+        rocket_model.set_id(rocket_id)
+        await self.put_rocket_by_id(rocket_id, rocket_model)
+        return
+
     @controller_exception_handler
     async def get_rocketpy_rocket_binary(self, rocket_id: str) -> bytes:
         """

src/models/flight.py

@@ -1,4 +1,7 @@
+import json
 from typing import Optional, Self, ClassVar, Literal
+
+from pydantic import BaseModel, Field, field_validator
 from src.models.interface import ApiBaseModel
 from src.models.rocket import RocketModel
 from src.models.environment import EnvironmentModel
@@ -69,3 +72,76 @@ def RETRIEVED(model_instance: type(Self)):
                 **model_instance.model_dump(),
             )
         )
+
+    @field_validator('environment', mode='before')
+    @classmethod
+    def _coerce_environment(cls, value):
+        if isinstance(value, str):
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError as exc:
+                raise ValueError(
+                    'Invalid JSON for environment payload'
+                ) from exc
+        return value
+
+    @field_validator('rocket', mode='before')
+    @classmethod
+    def _coerce_rocket(cls, value):
+        if isinstance(value, str):
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError as exc:
+                raise ValueError('Invalid JSON for rocket payload') from exc
+        return value
+
+
+class FlightPartialModel(BaseModel):
+    """Flight attributes required when rocket/environment are referenced."""
+
+    name: str = Field(default="flight")
+    rail_length: float = 1
+    time_overshoot: bool = True
+    terminate_on_apogee: bool = False
+    equations_of_motion: Literal['standard', 'solid_propulsion'] = 'standard'
+    inclination: float = 90.0
+    heading: float = 0.0
+    max_time: Optional[int] = None
+    max_time_step: Optional[float] = None
+    min_time_step: Optional[int] = None
+    rtol: Optional[float] = None
+    atol: Optional[float] = None
+    verbose: Optional[bool] = None
+
+    def assemble(
+        self,
+        *,
+        environment: EnvironmentModel,
+        rocket: RocketModel,
+    ) -> FlightModel:
+        """Compose a full flight model using referenced resources."""
+
+        flight_data = self.model_dump(exclude_none=True)
+        return FlightModel(
+            environment=environment,
+            rocket=rocket,
+            **flight_data,
+        )
+
+
+class FlightWithReferencesRequest(BaseModel):
+    """Payload for creating or updating flights via component references."""
+
+    environment_id: str
+    rocket_id: str
+    flight: FlightPartialModel
+
+    @field_validator('flight', mode='before')
+    @classmethod
+    def _coerce_flight(cls, value):
+        if isinstance(value, str):
+            try:
+                value = json.loads(value)
+            except json.JSONDecodeError as exc:
+                raise ValueError('Invalid JSON for flight payload') from exc
+        return value

src/models/motor.py

@@ -1,6 +1,7 @@
+import json
 from enum import Enum
 from typing import Optional, Tuple, List, Union, Self, ClassVar, Literal
-from pydantic import model_validator
+from pydantic import model_validator, field_validator
 
 from src.models.interface import ApiBaseModel
 from src.models.sub.tanks import MotorTank
@@ -57,6 +58,18 @@ class MotorModel(ApiBaseModel):
     ] = 'nozzle_to_combustion_chamber'
     reshape_thrust_curve: Union[bool, tuple] = False
 
+    @field_validator('tanks', mode='before')
+    @classmethod
+    def _coerce_tanks(cls, value):
+        if isinstance(value, str):
+            try:
+                value = json.loads(value)
+            except json.JSONDecodeError as exc:
+                raise ValueError('Invalid JSON for tanks payload') from exc
+        if isinstance(value, dict):
+            value = [value]
+        return value
+
     @model_validator(mode='after')
     # TODO: extend guard to check motor kinds and tank kinds specifics
     def validate_motor_kind(self):

src/models/rocket.py

@@ -1,4 +1,7 @@
+import json
 from typing import Optional, Tuple, List, Union, Self, ClassVar, Literal
+
+from pydantic import BaseModel, Field, field_validator
 from src.models.interface import ApiBaseModel
 from src.models.motor import MotorModel
 from src.models.sub.aerosurfaces import (
@@ -10,6 +13,15 @@
 )
 
 
+def _maybe_parse_json(value):
+    if isinstance(value, str):
+        try:
+            return json.loads(value)
+        except json.JSONDecodeError as exc:
+            raise ValueError('Invalid JSON payload') from exc
+    return value
+
+
 class RocketModel(ApiBaseModel):
     NAME: ClassVar = "rocket"
     METHODS: ClassVar = ("POST", "GET", "PUT", "DELETE")
@@ -37,6 +49,42 @@ class RocketModel(ApiBaseModel):
     rail_buttons: Optional[RailButtons] = None
     tail: Optional[Tail] = None
 
+    @field_validator('motor', mode='before')
+    @classmethod
+    def _coerce_motor(cls, value):
+        return _maybe_parse_json(value)
+
+    @field_validator('nose', mode='before')
+    @classmethod
+    def _coerce_nose(cls, value):
+        return _maybe_parse_json(value)
+
+    @field_validator('fins', mode='before')
+    @classmethod
+    def _coerce_fins(cls, value):
+        value = _maybe_parse_json(value)
+        if isinstance(value, dict):
+            value = [value]
+        return value
+
+    @field_validator('parachutes', mode='before')
+    @classmethod
+    def _coerce_parachutes(cls, value):
+        value = _maybe_parse_json(value)
+        if isinstance(value, dict):
+            value = [value]
+        return value
+
+    @field_validator('rail_buttons', mode='before')
+    @classmethod
+    def _coerce_rail_buttons(cls, value):
+        return _maybe_parse_json(value)
+
+    @field_validator('tail', mode='before')
+    @classmethod
+    def _coerce_tail(cls, value):
+        return _maybe_parse_json(value)
+
     @staticmethod
     def UPDATED():
         return
@@ -61,3 +109,53 @@ def RETRIEVED(model_instance: type(Self)):
                 **model_instance.model_dump(),
             )
         )
+
+
+class RocketPartialModel(BaseModel):
+    """Rocket attributes required when a motor is supplied by reference."""
+
+    radius: float
+    mass: float
+    motor_position: float
+    center_of_mass_without_motor: float
+    inertia: Union[
+        Tuple[float, float, float],
+        Tuple[float, float, float, float, float, float],
+    ] = (0, 0, 0)
+    power_off_drag: List[Tuple[float, float]] = Field(
+        default_factory=lambda: [(0, 0)]
+    )
+    power_on_drag: List[Tuple[float, float]] = Field(
+        default_factory=lambda: [(0, 0)]
+    )
+    coordinate_system_orientation: Literal['tail_to_nose', 'nose_to_tail'] = (
+        'tail_to_nose'
+    )
+    nose: NoseCone
+    fins: List[Fins]
+    parachutes: Optional[List[Parachute]] = None
+    rail_buttons: Optional[RailButtons] = None
+    tail: Optional[Tail] = None
+
+    def assemble(self, motor: MotorModel) -> RocketModel:
+        """Compose a full rocket model using the referenced motor."""
+
+        rocket_data = self.model_dump(exclude_none=True)
+        return RocketModel(motor=motor, **rocket_data)
+
+
+class RocketWithMotorReferenceRequest(BaseModel):
+    """Payload for creating or updating rockets via motor reference."""
+
+    motor_id: str
+    rocket: RocketPartialModel
+
+    @field_validator('rocket', mode='before')
+    @classmethod
+    def _coerce_rocket(cls, value):
+        if isinstance(value, str):
+            try:
+                value = json.loads(value)
+            except json.JSONDecodeError as exc:
+                raise ValueError('Invalid JSON for rocket payload') from exc
+        return value