Upload server/main.py with huggingface_hub

server/main.py

@@ -1,9 +1,16 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import Optional
+from typing import Optional, Any
 from server.environment import DataCentricEnvironment
 
-app = FastAPI(title="DataCentric-Env", version="0.1.0")
+app = FastAPI(
+    title="DataCentric-Env",
+    version="0.2.0",
+    description=(
+        "RL environment where an LLM acts as a data engineer. "
+        "Query specialist agents for recommendations, then apply them to fix a noisy dataset."
+    ),
+)
 env = DataCentricEnvironment()
 
 
@@ -11,6 +18,12 @@ class ResetRequest(BaseModel):
     difficulty: Optional[str] = None
 
 
+class ActionRequest(BaseModel):
+    action: str                         # query_cleaner | query_augmenter | ... | apply
+    rec_id: Optional[str] = None        # required for action="apply"
+    target_class: Optional[int] = None  # optional for query_augmenter
+
+
 @app.post("/reset")
 def reset(body: ResetRequest = None):
     difficulty = body.difficulty if body else None
@@ -18,8 +31,9 @@ def reset(body: ResetRequest = None):
 
 
 @app.post("/step")
-def step(action: dict):
-    return env.step(action)
+def step(body: ActionRequest):
+    action_dict = body.model_dump(exclude_none=True)
+    return env.step(action_dict)
 
 
 @app.get("/state")
@@ -29,4 +43,4 @@ def state():
 
 @app.get("/health")
 def health():
-    return {"status": "ok"}
+    return {"status": "ok", "version": "0.2.0"}