Spaces:
Sleeping
Sleeping
| import json | |
| from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect | |
| from pydantic import BaseModel | |
| from typing import Any, Optional | |
| try: | |
| from server.environment import DataCleaningEnvironment | |
| except ModuleNotFoundError: | |
| from environment import DataCleaningEnvironment | |
| app = FastAPI( | |
| title="Data Cleaning OpenEnv", | |
| description="A real-world data cleaning RL environment.", | |
| version="1.0.0", | |
| ) | |
| # ✅ Single shared instance — persists across requests | |
| env = DataCleaningEnvironment() | |
| class ResetRequest(BaseModel): | |
| task_name: str = "easy" | |
| class StepRequest(BaseModel): | |
| action_type: str | |
| column: str | |
| value: Optional[Any] = None | |
| def reset(request: Optional[ResetRequest] = None): | |
| try: | |
| task_name = request.task_name if request else "easy" | |
| return env.reset(task_name=task_name) | |
| except ValueError as e: | |
| raise HTTPException(status_code=400, detail=str(e)) | |
| def step(request: StepRequest): | |
| try: | |
| return env.step( | |
| action_type=request.action_type, | |
| column=request.column, | |
| value=request.value, | |
| ) | |
| except ValueError as e: | |
| raise HTTPException(status_code=400, detail=str(e)) | |
| def state(): | |
| return env.get_state() | |
| def health(): | |
| return {"status": "ok"} | |
| def list_tasks(): | |
| try: | |
| from tasks.graders import TASKS | |
| except ModuleNotFoundError: | |
| import sys, os | |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) | |
| from tasks.graders import TASKS | |
| return { | |
| k: { | |
| "name": v["name"], | |
| "description": v["description"], | |
| "difficulty": v["difficulty"], | |
| } | |
| for k, v in TASKS.items() | |
| } | |
| async def websocket_endpoint(websocket: WebSocket): | |
| await websocket.accept() | |
| ws_env = DataCleaningEnvironment() # isolated per connection | |
| try: | |
| while True: | |
| data = await websocket.receive_text() | |
| message = json.loads(data) | |
| msg_type = message.get("type") | |
| if msg_type == "reset": | |
| try: | |
| result = ws_env.reset(task_name=message.get("task_name", "easy")) | |
| await websocket.send_json({"status": "ok", **result}) | |
| except ValueError as e: | |
| await websocket.send_json({"status": "error", "detail": str(e)}) | |
| elif msg_type == "step": | |
| try: | |
| result = ws_env.step( | |
| action_type=message.get("action_type"), | |
| column=message.get("column"), | |
| value=message.get("value"), | |
| ) | |
| await websocket.send_json({"status": "ok", **result}) | |
| except ValueError as e: | |
| await websocket.send_json({"status": "error", "detail": str(e)}) | |
| elif msg_type == "state": | |
| await websocket.send_json({"status": "ok", **ws_env.get_state()}) | |
| elif msg_type == "health": | |
| await websocket.send_json({"status": "ok", "message": "healthy"}) | |
| else: | |
| await websocket.send_json({"status": "error", "detail": f"Unknown type: {msg_type}"}) | |
| except WebSocketDisconnect: | |
| pass | |
| def main(): | |
| """Main entry point for running the server.""" | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=7860) | |
| if __name__ == "__main__": | |
| main() |