Update main.py

main.py

@@ -47,7 +47,8 @@ app.add_middleware(
 
 # Single shared environment instance (stateful, per-session)
 _env = APIGatewayDefender()
-
+class ResetRequest(BaseModel):
+    task_id: str = "easy"
 
 # ─── Routes ──────────────────────────────────────────────────────────────────────
 
@@ -81,18 +82,13 @@ def root() -> Dict[str, Any]:
 
 
 @app.post("/reset")
-def reset(body: Dict[str, str] = None) -> Dict[str, Any]:
+def reset(req: ResetRequest = ResetRequest()) -> Dict[str, Any]:
     """
     Start a new episode.
-
-    Request body (JSON):
-        {"task_id": "easy" | "medium" | "hard"}
-
-    Returns the initial Observation.
+    Request body (JSON): {"task_id": "easy" | "medium" | "hard"}
     """
-    task_id = (body or {}).get("task_id", "easy")
     try:
-        obs: Observation = _env.reset(task_id=task_id)
+        obs: Observation = _env.reset(task_id=req.task_id)
     except ValueError as exc:
         raise HTTPException(status_code=422, detail=str(exc))
     return obs.model_dump()