server/routes.py

import os

from fastapi import APIRouter, HTTPException

from env.models import Action
from server.app import env

router = APIRouter()


@router.post("/reset")
def reset() -> dict:
    return env.reset().model_dump()


@router.post("/step")
def step(action: Action) -> dict:
    try:
        obs, reward, done, info = env.step(action)
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    return {
        "observation": obs.model_dump(),
        "reward": reward,
        "done": done,
        "info": info,
    }


@router.get("/state")
def state() -> dict:
    return env.state()