server/app.py

"""
FastAPI application for the Data Cleaning Env Environment.

Endpoints provided by openenv create_app():
    POST /reset    – Reset environment
    POST /step     – Execute an action
    GET  /state    – Current environment state
    GET  /schema   – Action/Observation JSON schemas
    WS   /ws       – WebSocket endpoint

Additional hackathon-required endpoints:
    GET  /health   – Health check
    GET  /tasks    – List tasks + action schema for each difficulty
    POST /grader   – Return grader score for current episode
    POST /baseline – Run deterministic baseline agent on all 3 tasks
"""

try:
    from openenv.core.env_server.http_server import create_app
except Exception as e:
    raise ImportError(
        "openenv is required. Install dependencies with: uv sync"
    ) from e

from models import DataCleaningAction, DataCleaningObservation
from server.data_cleaning_env_environment import DataCleaningEnvironment

from fastapi import HTTPException
from fastapi.responses import JSONResponse

# ---------------------------------------------------------------------------
# Base OpenEnv application
# ---------------------------------------------------------------------------
app = create_app(
    DataCleaningEnvironment,
    DataCleaningAction,
    DataCleaningObservation,
    env_name="data_cleaning_env",
    max_concurrent_envs=1,
)

# ---------------------------------------------------------------------------
# Shared environment instance for grader/baseline
# ---------------------------------------------------------------------------
_env: DataCleaningEnvironment | None = None

def _get_env() -> DataCleaningEnvironment:
    global _env
    if _env is None:
        _env = DataCleaningEnvironment()
    return _env

# ---------------------------------------------------------------------------
# GET /health
# ---------------------------------------------------------------------------
@app.get("/health")
def health():
    """Simple health check endpoint."""
    return {"status": "ok"}

# ---------------------------------------------------------------------------
# GET /tasks
# ---------------------------------------------------------------------------
@app.get("/tasks")
def get_tasks():
    """Return list of available tasks."""
    try:
        tasks = DataCleaningEnvironment.tasks()
        return JSONResponse(content={"tasks": tasks})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ---------------------------------------------------------------------------
# POST /grader
# ---------------------------------------------------------------------------
@app.post("/grader")
def run_grader():
    """
    Score the current episode.

    Must return:
        {"score": float between 0 and 1}
    """
    try:
        result = _get_env().grade()

        if isinstance(result, dict) and "score" in result:
            return JSONResponse(content=result)
        return JSONResponse(content={"score": float(result)})

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ---------------------------------------------------------------------------
# POST /baseline
# ---------------------------------------------------------------------------
@app.post("/baseline")
def run_baseline():
    """
    Run deterministic baseline agent across all tasks.
    """
    try:
        result = _get_env().run_baseline()
        return JSONResponse(content=result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def main(host: str = "0.0.0.0", port: int = 8000) -> None:
    import uvicorn
    uvicorn.run("server.app:app", host=host, port=port)

if __name__ == "__main__":
    main()