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 with: uv sync") from e

# Robust imports that work both locally and inside Docker
try:
    from ..models import DataCleaningAction, DataCleaningObservation
    from .data_cleaning_env_environment import DataCleaningEnvironment
except ImportError:
    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 (handles /reset, /step, /state, /schema, /ws)
# ---------------------------------------------------------------------------
app = create_app(
    DataCleaningEnvironment,
    DataCleaningAction,
    DataCleaningObservation,
    env_name="data_cleaning_env",
    max_concurrent_envs=1,
)

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

def _get_env() -> DataCleaningEnvironment:
    global _env
    if _env is None:
        _env = DataCleaningEnvironment()
        _env.reset(task="easy")
    return _env

# ---------------------------------------------------------------------------
# GET /health
# ---------------------------------------------------------------------------
@app.get("/health")
def health():
    """Health check — required by HF Space ping."""
    return {"status": "ok"}


# ---------------------------------------------------------------------------
# GET /tasks  — required by hackathon checklist
# ---------------------------------------------------------------------------
@app.get("/tasks")
def get_tasks():
    """Return list of available tasks with action schemas."""
    try:
        return JSONResponse(content={"tasks": DataCleaningEnvironment.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. Returns score in [0, 1]."""
    try:
        result = _get_env().grade()
        return JSONResponse(content=result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ---------------------------------------------------------------------------
# POST /baseline  — required by hackathon checklist
# ---------------------------------------------------------------------------
@app.post("/baseline")
def run_baseline():
    """Run deterministic rule-based baseline on all 3 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(app, host=host, port=port)   


if __name__ == "__main__":
    main()