Spaces:
Sleeping
Sleeping
File size: 3,579 Bytes
fee2f22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | 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
@app.post("/reset")
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))
@app.post("/step")
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))
@app.get("/state")
def state():
return env.get_state()
@app.get("/health")
def health():
return {"status": "ok"}
@app.get("/tasks")
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()
}
@app.websocket("/ws")
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() |