src/auditenv/server.py

from __future__ import annotations

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from auditenv.models import AuditAction, AuditObservation, EnvState, StepResult, TaskId
from auditenv.state import AuditEnvRuntime


app = FastAPI(title="AuditEnv", version="0.1.0")
runtime = AuditEnvRuntime(default_seed=42)


class ResetRequest(BaseModel):
    task_id: TaskId = "easy"
    seed: int = 42


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/reset", response_model=AuditObservation)
def reset(req: ResetRequest) -> AuditObservation:
    return runtime.reset(task_id=req.task_id, seed=req.seed)


@app.post("/step", response_model=StepResult)
def step(action: AuditAction) -> StepResult:
    try:
        return runtime.step(action)
    except RuntimeError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except ValueError as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc


@app.get("/state", response_model=EnvState)
def state() -> EnvState:
    try:
        return runtime.state()
    except RuntimeError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc