EC256/openenv-data-engineering

1"""2FastAPI server for the OpenEnv environment.3Exposes /health, /reset, /step endpoints and runs on port 7860.4"""5 6from fastapi import FastAPI, HTTPException7from fastapi.responses import JSONResponse8 9app = FastAPI(title="OpenEnv: Advanced Data Engineering Pipeline")10 11# Import environment components12from environment import DataEnv13from tasks import TASKS14 15# Store environment instance globally16env = DataEnv()17 18 19@app.get("/health")20async def health():21    """Health check endpoint."""22    return JSONResponse(23        status_code=200,24        content={25            "status": "healthy",26            "service": "openenv-data-engineering-pipeline",27            "version": "2.0.0",28        },29    )30 31 32@app.get("/")33async def root():34    """Root endpoint."""35    return JSONResponse(36        status_code=200,37        content={38            "message": "OpenEnv: Advanced Multi-Table Data Engineering Pipeline",39            "docs": "/docs",40            "health": "/health",41        },42    )43 44 45@app.post("/reset")46async def reset_environment(task_id: str = "easy_data_cleaning"):47    """48    Reset the environment to the initial state.49    Returns observation AND info (Gymnasium protocol).50    """51    try:52        obs = env.reset(task_id=task_id)53        # CRITICAL: Judge expects both observation and info54        return JSONResponse(55            status_code=200,56            content={57                "observation": {58                    "dataset_sample": obs.dataset_sample,59                    "columns": obs.columns,60                    "dtypes": obs.dtypes,61                    "total_rows": obs.total_rows,62                    "missing_values": obs.missing_values,63                    "task_description": obs.task_description,64                    "feedback": obs.feedback,65                    "debug_hints": obs.debug_hints,66                },67                "info": {},68            },69        )70    except Exception as e:71        raise HTTPException(status_code=500, detail=str(e))72 73 74@app.post("/step")75async def step_environment(action: dict):76    """77    Execute an action and return observation, reward, done, info.78    """79    try:80 81        class DummyAction:82            def __init__(self, d):83                for k, v in d.items():84                    setattr(self, k, v)85                if not hasattr(self, "action_type"):86                    self.action_type = "submit"87 88        action_obj = DummyAction(action)89        obs, reward, done, info = env.step(action_obj)90 91        return JSONResponse(92            status_code=200,93            content={94                "observation": {95                    "dataset_sample": obs.dataset_sample,96                    "columns": obs.columns,97                    "dtypes": obs.dtypes,98                    "total_rows": obs.total_rows,99                    "missing_values": obs.missing_values,100                    "task_description": obs.task_description,101                    "feedback": obs.feedback,102                    "debug_hints": obs.debug_hints,103                },104                "reward": {105                    "value": reward.value,106                    "message": reward.message,107                },108                "done": done,109                "info": info,110            },111        )112    except Exception as e:113        raise HTTPException(status_code=500, detail=str(e))114 115 116@app.get("/tasks")117async def get_tasks():118    """Return available tasks."""119    return JSONResponse(120        status_code=200,121        content={122            "tasks": [123                {124                    "id": task_id,125                    "description": TASKS[task_id].description,126                }127                for task_id in TASKS.keys()128            ]129        },130    )131 132 133if __name__ == "__main__":134    import uvicorn135 136    uvicorn.run(app, host="0.0.0.0", port=7860)137 138 139def main():140    """Entry point for the server."""141    import uvicorn142 143    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)144