server/app.py

"""
FastAPI application exposing the OpenEnv-compatible HTTP API.

Endpoints:
  GET  /health       Health check
  GET  /metadata     Environment info
  GET  /schema       Action / observation / state schemas
  POST /reset        Start new episode
  POST /step         Execute cleaning action (with 30s timeout)
  GET  /state        Episode metadata
  POST /state        Episode metadata (backward compat)
  GET  /profile      Rich data quality profile of current DataFrame
  GET  /report       Full episode cleaning summary (health certificate)
  GET  /export       Download current cleaned DataFrame as CSV
"""

import asyncio
import os
from typing import Any, Dict, Optional
from fastapi import Body, FastAPI, HTTPException
from fastapi.responses import PlainTextResponse, HTMLResponse
from pydantic import BaseModel
import uvicorn

from models import DataCleaningAction, DataCleaningObservation, DataCleaningState, EpisodeReport
from server.environment import DataCleaningEnvironment

app = FastAPI(
    title="Data Cleaning OpenEnv",
    description=(
        "A real-world data cleaning environment for AI agent training and evaluation. "
        "An agent interacts with dirty pandas DataFrames through a standard reset/step/state API, "
        "learning to fix missing values, duplicates, format inconsistencies, outliers, and dtype errors. "
        "Grounded in CleanAgent (2024), AutoDCWorkflow (EMNLP 2025), and Meta-scale data quality principles."
    ),
    version="0.2.0",
)

# Single shared environment instance
env = DataCleaningEnvironment()

STEP_TIMEOUT_SECONDS = 30


class ResetRequest(BaseModel):
    task_id: Optional[int] = None


class StepResponse(BaseModel):
    observation: DataCleaningObservation
    reward: float
    done: bool
    info: dict = {}


# ------------------------------------------------------------------
# Core OpenEnv routes
# ------------------------------------------------------------------

@app.get("/", response_class=HTMLResponse, include_in_schema=False)
def ui():
    """DataMedic — live agent monitoring dashboard."""
    ui_path = os.path.join(os.path.dirname(__file__), "ui.html")
    with open(ui_path, "r") as f:
        return HTMLResponse(content=f.read())


@app.get("/health")
def health():
    return {"status": "healthy"}


@app.get("/metadata")
def metadata():
    return {
        "name": "data-cleaning-env",
        "description": (
            "A real-world data cleaning RL environment. The agent diagnoses dirty datasets, "
            "plans a treatment, executes cleaning operations step-by-step, and produces a "
            "health certificate — grounded in AutoDCWorkflow, CleanAgent, and HoloClean research."
        ),
        "version": "0.2.0",
        "tags": ["openenv", "data-cleaning", "rl", "real-world", "agentic"],
        "tasks": [
            {"id": "task1", "name": "Fill Missing Values",               "difficulty": "easy"},
            {"id": "task2", "name": "Fix Formats and Remove Duplicates", "difficulty": "medium"},
            {"id": "task3", "name": "Full Cleaning Pipeline",            "difficulty": "hard"},
            {"id": "task4", "name": "Multi-Source Schema Alignment + Merge", "difficulty": "expert"},
        ],
        "observation_extras": ["dq_metrics", "tried_operations", "plan"],
        "papers": [
            "Bendinelli et al. 2025 — LLM Agents for Cleaning Tabular ML Datasets (arXiv:2503.06664)",
            "CleanAgent — Qi & Wang 2024 (arXiv:2403.08291)",
            "AutoDCWorkflow — EMNLP 2025 Findings",
            "HoloClean — Rekatsinas et al. 2017",
        ],
    }


@app.get("/schema")
def schema():
    return {
        "action": {
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": [
                        "fill_missing", "drop_duplicates", "fix_format",
                        "replace_value", "drop_outliers", "fix_dtype",
                        "align_schema", "merge_sources",
                    ],
                },
                "column": {"type": "string", "nullable": True},
                "params":  {"type": "object", "nullable": True},
            },
            "required": ["operation"],
        },
        "observation": {
            "type": "object",
            "properties": {
                "done":             {"type": "boolean"},
                "reward":           {"type": "number"},
                "data_preview":     {"type": "string"},
                "data_shape":       {"type": "array", "items": {"type": "integer"}},
                "missing_counts":   {"type": "object"},
                "duplicate_count":  {"type": "integer"},
                "dtype_issues":     {"type": "object"},
                "task_description": {"type": "string"},
                "message":          {"type": "string"},
                "step_count":       {"type": "integer"},
                "current_score":    {"type": "number"},
                "dq_metrics":       {"type": "object", "description": "Completeness/uniqueness/validity %"},
                "tried_operations": {"type": "array",  "description": "Operations already applied"},
                "plan":             {"type": "array",  "description": "Agent recommended next actions"},
            },
        },
        "state": {
            "type": "object",
            "properties": {
                "episode_id":       {"type": "string"},
                "task_id":          {"type": "integer"},
                "step_count":       {"type": "integer"},
                "max_steps":        {"type": "integer"},
                "total_errors":     {"type": "integer"},
                "errors_remaining": {"type": "integer"},
            },
        },
    }


@app.post("/reset", response_model=StepResponse)
def reset(req: ResetRequest = ResetRequest()):
    try:
        obs = env.reset(task_id=req.task_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    return StepResponse(observation=obs, reward=obs.reward, done=False)


@app.post("/step", response_model=StepResponse)
async def step(body: Dict[str, Any] = Body(...)):
    """
    Accept both openenv-core wrapped format:
        {"action": {"operation": "...", ...}, "timeout_s": 15}
    and direct format:
        {"operation": "...", "column": "...", "params": {...}}
    Times out after 30 seconds to prevent hanging during evaluation.
    """
    action_data = body.get("action", body)
    try:
        action = DataCleaningAction(**action_data)
        loop = asyncio.get_event_loop()
        obs = await asyncio.wait_for(
            loop.run_in_executor(None, env.step, action),
            timeout=STEP_TIMEOUT_SECONDS,
        )
    except asyncio.TimeoutError:
        raise HTTPException(status_code=504, detail=f"Step timed out after {STEP_TIMEOUT_SECONDS}s")
    except (TypeError, KeyError, Exception) as e:
        raise HTTPException(status_code=400, detail=str(e))
    return StepResponse(observation=obs, reward=obs.reward, done=obs.done)


@app.get("/state", response_model=DataCleaningState)
def state_get():
    return env.state()


@app.post("/state", response_model=DataCleaningState)
def state_post():
    return env.state()


# ------------------------------------------------------------------
# Phase 2: Intelligence endpoints
# ------------------------------------------------------------------

@app.get("/profile")
def profile():
    """
    Rich data quality profile of the current DataFrame.

    Returns per-column statistics (null %, unique %, min/max/mean for numerics,
    top values for categoricals) plus dataset-level DQ metrics:
    completeness %, uniqueness %, validity %.

    Inspired by standard Data Quality dimensions (ISO 8000) and
    Meta's data schematization approach.
    """
    try:
        return env.get_profile()
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/report", response_model=EpisodeReport)
def report():
    """
    Full episode cleaning summary — the 'health certificate'.

    Returns: initial vs final score, score improvement, step efficiency,
    ordered list of operations applied, issues fixed by category,
    and final DQ metrics. Call after the episode completes for best results.
    """
    try:
        return env.get_report()
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/export")
def export():
    """
    Download the current (cleaned) DataFrame as a CSV file.

    Returns the live state of the DataFrame — call after the agent
    finishes cleaning to get the cleaned output.
    """
    try:
        csv_data = env.get_export()
        return PlainTextResponse(
            content=csv_data,
            media_type="text/csv",
            headers={"Content-Disposition": "attachment; filename=cleaned_data.csv"},
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


# ------------------------------------------------------------------
# Entry point
# ------------------------------------------------------------------

def main():
    uvicorn.run("server.app:app", host="0.0.0.0", port=8000, workers=1)


if __name__ == "__main__":
    main()