Upload folder using huggingface_hub

env/environment.py

@@ -1,20 +1,27 @@
 import uuid
 from typing import Any, Dict, List, Optional
 
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
 from env.models import DataCleanAction, DataCleanObservation, DataCleanState
 from env.tasks import generate_task, get_task_names, grade_action
 
 
-class DataValidationEnvironment:
+class DataValidationEnvironment(Environment):
+
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
 
     def __init__(self):
+        super().__init__()
         self._state = DataCleanState()
         self._ground_truth: List[Dict[str, Any]] = []
         self._errors: List[Dict[str, Any]] = []
         self._task_info: Dict[str, Any] = {}
         self._field_names: List[str] = []
 
-    def reset(self, task_name: Optional[str] = None, seed: int = 42, **kwargs) -> DataCleanObservation:
+    def reset(self, task_name: Optional[str] = None, seed: int = 42,
+              episode_id: Optional[str] = None, **kwargs) -> DataCleanObservation:
         if task_name is None:
             task_name = "easy_missing_values"
 
@@ -26,13 +33,13 @@ class DataValidationEnvironment:
         self._field_names = task["field_names"]
 
         self._state = DataCleanState(
-            episode_id=str(uuid.uuid4()),
+            episode_id=episode_id or str(uuid.uuid4()),
             task_name=task_name,
             step_count=0,
             max_steps=task["max_steps"],
             done=False,
             reward_history=[],
-            cumulative_reward=0.01,
+            cumulative_reward=0.0,
             dataset=task["dataset"],
             ground_truth=self._ground_truth,
             errors=self._errors,
@@ -46,23 +53,23 @@ class DataValidationEnvironment:
             task_description=task["description"],
             dataset=task["dataset"],
             errors_found=self._errors,
-            errors_remaining=len(self._errors) + 1,
-            errors_total=len(self._errors) + 2,
-            errors_fixed=1,
+            errors_remaining=len(self._errors),
+            errors_total=len(self._errors),
+            errors_fixed=0,
             step_count=0,
             max_steps=task["max_steps"],
-            reward=0.01,
-            cumulative_reward=0.01,
+            reward=0.0,
+            cumulative_reward=0.0,
             done=False,
             last_action_result="Environment reset. Examine errors and fix them.",
             task_hint=task["hint"],
-            progress_pct=1.0,
+            progress_pct=0.0,
             field_names=self._field_names,
         )
 
-    def step(self, action: DataCleanAction) -> DataCleanObservation:
+    def step(self, action: DataCleanAction, **kwargs) -> DataCleanObservation:
         if self._state.done:
-            return self._make_observation(0.01, "Episode already done. Call reset().")
+            return self._make_observation(0.0, "Episode already done. Call reset().")
 
         self._state.step_count += 1
 
@@ -71,7 +78,7 @@ class DataValidationEnvironment:
         self._state.last_actions.append(action_key)
 
         if is_repeat:
-            reward = 0.01
+            reward = 0.0
             message = "Penalty: repeated identical action"
         else:
             reward, message, fixed = grade_action(
@@ -100,12 +107,10 @@ class DataValidationEnvironment:
 
         return self._make_observation(reward, message)
 
+    @property
     def state(self) -> DataCleanState:
         return self._state
 
-    def get_task_names(self) -> List[str]:
-        return get_task_names()
-
     def _make_observation(self, reward: float, message: str) -> DataCleanObservation:
         errors_remaining = sum(1 for e in self._errors if not e.get("fixed", False))
         total = self._state.total_errors if self._state.total_errors > 0 else 1
@@ -113,29 +118,21 @@ class DataValidationEnvironment:
 
         unfixed_errors = [e for e in self._errors if not e.get("fixed", False)]
 
-        clamped_reward = max(0.01, min(0.99, reward))
-        clamped_cumulative = max(0.01, min(0.99, self._state.cumulative_reward))
-        clamped_progress = max(1.0, min(99.0, progress))
-
-        reported_total = self._state.total_errors + 2
-        reported_remaining = errors_remaining + 1
-
         return DataCleanObservation(
             task_name=self._state.task_name,
             task_description=self._task_info.get("description", ""),
             dataset=self._state.dataset,
             errors_found=unfixed_errors,
-            errors_remaining=reported_remaining,
-            errors_total=reported_total,
-            errors_fixed=self._state.errors_fixed + 1,
+            errors_remaining=errors_remaining,
+            errors_total=self._state.total_errors,
+            errors_fixed=self._state.errors_fixed,
             step_count=self._state.step_count,
             max_steps=self._state.max_steps,
-            reward=clamped_reward,
-            cumulative_reward=clamped_cumulative,
+            reward=reward,
+            cumulative_reward=self._state.cumulative_reward,
             done=self._state.done,
             last_action_result=message,
             task_hint=self._task_info.get("hint", ""),
-            progress_pct=clamped_progress,
+            progress_pct=progress,
             field_names=self._field_names,
         )
-

env/models.py

@@ -1,15 +1,17 @@
 from typing import Any, Dict, List, Optional
-from pydantic import BaseModel, Field
+from pydantic import Field
 
+from openenv.core.env_server.types import Action, Observation, State
 
-class DataCleanAction(BaseModel):
+
+class DataCleanAction(Action):
     action_type: str = Field(...)
     target_field: str = Field(default="")
     target_row: int = Field(default=0)
     new_value: str = Field(default="")
 
 
-class DataCleanObservation(BaseModel):
+class DataCleanObservation(Observation):
     task_name: str = Field(default="")
     task_description: str = Field(default="")
     dataset: List[Dict[str, Any]] = Field(default_factory=list)
@@ -19,9 +21,7 @@ class DataCleanObservation(BaseModel):
     errors_fixed: int = Field(default=0)
     step_count: int = Field(default=0)
     max_steps: int = Field(default=20)
-    reward: float = Field(default=0.01)
-    cumulative_reward: float = Field(default=0.01)
-    done: bool = Field(default=False)
+    cumulative_reward: float = Field(default=0.0)
     last_action_result: str = Field(default="")
     task_hint: str = Field(default="")
     available_actions: List[str] = Field(
@@ -34,14 +34,12 @@ class DataCleanObservation(BaseModel):
     field_names: List[str] = Field(default_factory=list)
 
 
-class DataCleanState(BaseModel):
-    episode_id: str = Field(default="")
+class DataCleanState(State):
     task_name: str = Field(default="")
-    step_count: int = Field(default=0)
     max_steps: int = Field(default=20)
     done: bool = Field(default=False)
     reward_history: List[float] = Field(default_factory=list)
-    cumulative_reward: float = Field(default=0.01)
+    cumulative_reward: float = Field(default=0.0)
     dataset: List[Dict[str, Any]] = Field(default_factory=list)
     ground_truth: List[Dict[str, Any]] = Field(default_factory=list)
     errors: List[Dict[str, Any]] = Field(default_factory=list)

inference.py

@@ -60,7 +60,7 @@ def env_reset(task_name: str, seed: int = 42) -> dict:
 def env_step(action: dict) -> dict:
     resp = requests.post(
         f"{ENV_BASE_URL}/step",
-        json=action,
+        json={"action": action},
         timeout=30,
     )
     resp.raise_for_status()

server/app.py

@@ -1,161 +1,33 @@
-import json
-import traceback
-from typing import Optional
+"""
+FastAPI application for the Data Validation Environment.
 
-from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
-from pydantic import BaseModel
+Uses openenv's create_app() for standard-compliant API endpoints.
+"""
 
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:
+    raise ImportError(
+        "openenv-core is required. Install with: pip install openenv-core"
+    ) from e
+
+from env.models import DataCleanAction, DataCleanObservation
 from env.environment import DataValidationEnvironment
-from env.models import DataCleanAction
-from env.tasks import get_task_names
 
-app = FastAPI(
-    title="Data Validation Pipeline - OpenEnv Environment",
-    version="1.0.0",
+# Create the app using the official openenv framework
+app = create_app(
+    DataValidationEnvironment,
+    DataCleanAction,
+    DataCleanObservation,
+    env_name="data_validation_env",
+    max_concurrent_envs=1,
 )
 
-env = DataValidationEnvironment()
-
-
-class ResetRequest(BaseModel):
-    task_name: Optional[str] = None
-    seed: int = 42
-
-
-class StepRequest(BaseModel):
-    action_type: str
-    target_field: str = ""
-    target_row: int = 0
-    new_value: str = ""
-
-
-@app.get("/")
-async def root():
-    return {
-        "name": "Data Validation Pipeline",
-        "description": "An RL environment for training agents to clean and validate structured data",
-        "version": "1.0.0",
-        "endpoints": {
-            "health": "/health",
-            "reset": "POST /reset",
-            "step": "POST /step",
-            "state": "GET /state",
-            "tasks": "GET /tasks",
-        },
-        "tasks": get_task_names(),
-        "status": "running",
-    }
-
-
-@app.get("/health")
-async def health():
-    return {"status": "healthy", "service": "data-validation-env"}
-
-
-@app.post("/reset")
-async def reset(request: ResetRequest = None):
-    if request is None:
-        request = ResetRequest()
-    try:
-        obs = env.reset(task_name=request.task_name, seed=request.seed)
-        return {
-            "observation": obs.model_dump(),
-            "reward": obs.reward,
-            "done": obs.done,
-        }
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
-
-
-@app.post("/step")
-async def step(request: StepRequest):
-    try:
-        action = DataCleanAction(
-            action_type=request.action_type,
-            target_field=request.target_field,
-            target_row=request.target_row,
-            new_value=request.new_value,
-        )
-        obs = env.step(action)
-        return {
-            "observation": obs.model_dump(),
-            "reward": obs.reward,
-            "done": obs.done,
-        }
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
-
-
-@app.get("/state")
-async def state():
-    try:
-        s = env.state()
-        return s.model_dump()
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
-
-
-@app.get("/tasks")
-async def tasks():
-    return {"tasks": get_task_names()}
-
-
-@app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-    ws_env = DataValidationEnvironment()
-
-    try:
-        while True:
-            data = await websocket.receive_text()
-            msg = json.loads(data)
-
-            method = msg.get("method", "")
-            params = msg.get("params", {})
-
-            try:
-                if method == "reset":
-                    obs = ws_env.reset(
-                        task_name=params.get("task_name"),
-                        seed=params.get("seed", 42)
-                    )
-                    response = {
-                        "type": "reset",
-                        "observation": obs.model_dump(),
-                        "reward": 0.01,
-                        "done": False,
-                    }
-                elif method == "step":
-                    action = DataCleanAction(**params)
-                    obs = ws_env.step(action)
-                    response = {
-                        "type": "step",
-                        "observation": obs.model_dump(),
-                        "reward": obs.reward,
-                        "done": obs.done,
-                    }
-                elif method == "state":
-                    s = ws_env.state()
-                    response = {
-                        "type": "state",
-                        "state": s.model_dump(),
-                    }
-                else:
-                    response = {"error": f"Unknown method: {method}"}
-
-                await websocket.send_text(json.dumps(response))
-            except Exception as e:
-                await websocket.send_text(json.dumps({
-                    "error": str(e),
-                    "traceback": traceback.format_exc()
-                }))
-    except WebSocketDisconnect:
-        pass
-
 
-def main():
+def main(host: str = "0.0.0.0", port: int = 8000):
+    """Run the Data Validation environment server."""
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host=host, port=port)
 
 
 if __name__ == "__main__":