Upload 22 files

New space created, with the updated code

.dockerignore

@@ -0,0 +1,10 @@
+.env
+.env.local
+.env.*.local
+conda-env/
+.venv/
+venv/
+__pycache__/
+.pytest_cache/
+*.pyc
+.DS_Store

.env.example

@@ -0,0 +1,11 @@
+API_BASE_URL=https://api.groq.com/openai/v1
+MODEL_NAME=meta-llama/llama-4-scout-17b-16e-instruct
+HF_TOKEN=
+# Available TASK_ID values:
+# task1_easy
+# task2_medium
+# task3_hard
+# task4_medium_alt
+# task5_hard_alt
+TASK_ID=task1_easy
+MAX_STEPS=15

Dockerfile

@@ -0,0 +1,19 @@
+FROM python:3.11-slim
+
+# Non-root user required by Hugging Face Spaces
+RUN useradd -m -u 1000 appuser
+
+WORKDIR /app
+
+# Install dependencies first (better layer caching)
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY --chown=appuser:appuser . .
+
+USER appuser
+
+EXPOSE 7860
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,12 +1,104 @@
 ---
-title: Openenv
-emoji: 🏃
-colorFrom: red
+title: Data Cleaning OpenEnv Benchmark
+emoji: 🧹
+colorFrom: blue
 colorTo: green
 sdk: docker
 pinned: false
-license: mit
-short_description: Data Cleaning OpenEnv Benchmark
+tags:
+  - openenv
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Data Cleaning OpenEnv Benchmark
+
+A practical benchmark where LLM agents clean messy tabular datasets through a structured action API.
+
+## Why This Matters
+
+Data cleaning still takes a large share of real analytics work. This environment tests whether an agent can detect and correct common data quality problems such as duplicates, missing values, inconsistent formats, and outliers.
+
+## Tasks
+
+| ID | Difficulty | Description |
+|----|-----------|-------------|
+| `task1_easy` | Easy | Remove exact duplicates, fill missing emails and ages, standardise country names |
+| `task2_medium` | Medium | Normalise mixed date formats, convert price strings to float, fix category typos |
+| `task3_hard` | Hard | Resolve duplicate user IDs, clip session outliers, fix invalid bounce rates |
+| `task4_medium_alt` | Medium | Alternate order-cleaning scenario that uses the same grader contract as `task2_medium` |
+| `task5_hard_alt` | Hard | Alternate analytics-cleaning scenario that uses the same grader contract as `task3_hard` |
+
+Each task is graded independently, and scores are always strictly between 0 and 1.
+
+## Action Space
+
+| Action | Required Fields |
+|--------|----------------|
+| `fill_missing` | `column`, `strategy` (`mean`/`median`/`mode`/`constant`), `value` when needed |
+| `standardize_values` | `column`, `mapping` |
+| `remove_duplicates` | None |
+| `remove_row` | `row_id` |
+| `convert_type` | `column`, `target_type` |
+| `clip_outliers` | `column`, `lower`, `upper` |
+| `submit` | None |
+
+## Observation Space
+
+Each step the agent receives `table_preview`, `schema_info`, `issues_detected`, `cleaning_log`, `valid_actions`, `step`, and `max_steps`.
+
+## Reward Design
+
+Correct cleaning actions receive positive intermediate rewards, wasted actions receive small penalties, invalid actions receive larger penalties, and `submit` returns the final grader score.
+
+## Setup & Local Run
+
+```bash
+git clone https://huggingface.co/spaces/AnkushRaheja/data-cleaning-benchmark
+cd data-cleaning-benchmark
+pip install -r requirements.txt
+uvicorn app:app --port 7860
+```
+
+## Run Baseline
+
+```bash
+export API_BASE_URL="https://api.groq.com/openai/v1"
+export MODEL_NAME="meta-llama/llama-4-scout-17b-16e-instruct"
+export HF_TOKEN="$GROQ_API_KEY"
+export TASK_ID="task1_easy"
+python inference.py
+```
+
+## Docker
+
+```bash
+docker build -t data-cleaning-benchmark .
+docker run -p 7860:7860 \
+  -e API_BASE_URL="https://api.groq.com/openai/v1" \
+  -e MODEL_NAME="meta-llama/llama-4-scout-17b-16e-instruct" \
+  -e HF_TOKEN="$GROQ_API_KEY" \
+  data-cleaning-benchmark
+```
+
+## Baseline Scores
+
+| Task | Score |
+|------|-------|
+| task1_easy | 0.99 |
+| task2_medium | 0.99 |
+| task3_hard | 0.97 |
+| task4_medium_alt | 0.99 |
+| task5_hard_alt | 0.97 |
+
+## API Reference
+
+| Method | Endpoint | Description |
+|--------|----------|-------------|
+| GET | `/health` | Health check |
+| POST | `/reset` | Start new episode `{"task_id": "task1_easy"}` |
+| POST | `/step` | Submit action and receive reward (compat route with `session_id` in body/query) |
+| POST | `/step/{session_id}` | Legacy route for direct session addressing |
+| GET | `/state` | Retrieve state by query (`session_id`) |
+| GET | `/state/{session_id}` | Legacy route for direct session addressing |
+| GET | `/tasks` | List all tasks |
+| GET | `/metadata` | Benchmark metadata including task and score-range contract |
+| GET | `/schema` | JSON schemas for action/observation/step response |

app.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import uuid
+from typing import Any, Dict, Optional
+
+from fastapi import FastAPI, HTTPException, Query
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+
+from env.environment import DataCleaningEnv
+from env.models import Action, Observation, StepResult
+from env.tasks import list_tasks as list_task_specs
+
+app = FastAPI(
+    title="Data Cleaning OpenEnv Benchmark",
+    version="1.0.0",
+    description="LLM agent benchmark for real-world data cleaning tasks.",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+sessions: Dict[str, DataCleaningEnv] = {}
+
+
+@app.get("/")
+def root():
+    tasks = list_task_specs()
+    return {
+        "name": "Data Cleaning OpenEnv Benchmark",
+        "version": "1.0.0",
+        "tasks": tasks,
+        "api": {
+            "reset": "POST /reset",
+            "step": "POST /step/{session_id}",
+            "step_compat": "POST /step",
+            "state": "GET  /state/{session_id}",
+            "state_compat": "GET  /state?session_id=...",
+            "metadata": "GET  /metadata",
+            "schema": "GET  /schema",
+            "mcp": "GET|POST /mcp",
+            "health": "GET  /health",
+        },
+    }
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "sessions_active": len(sessions)}
+
+
+class ResetRequest(BaseModel):
+    task_id: Optional[str] = None
+
+
+@app.post("/reset")
+def reset(body: ResetRequest = ResetRequest()):
+    session_id = str(uuid.uuid4())
+    env = DataCleaningEnv()
+    obs = env.reset(task_id=body.task_id)
+    sessions[session_id] = env
+    return {
+        "session_id": session_id,
+        "observation": obs.model_dump(),
+        "reward": 0.0,
+        "done": False,
+        "info": {
+            "error": None,
+            "cumulative_reward": env.cumulative_reward,
+            "raw_cumulative_reward": env.raw_cumulative_reward,
+            "final_score": env.final_score,
+            "step": env.step_count,
+        },
+    }
+
+
+@app.post("/step")
+def step_compat(
+    payload: Dict[str, Any],
+    session_id: Optional[str] = Query(default=None),
+):
+    payload_session_id = payload.get("session_id")
+    resolved_session_id = _resolve_session_id(payload_session_id or session_id)
+    action_payload = payload.get("action", payload)
+
+    if not isinstance(action_payload, dict):
+        raise HTTPException(status_code=400, detail="Action payload must be an object")
+    if "type" not in action_payload:
+        raise HTTPException(status_code=400, detail="Action payload requires 'type'")
+
+    action = Action(**action_payload)
+    env = _get_session(resolved_session_id)
+    result = env.step(action)
+    return result.model_dump()
+
+
+@app.post("/step/{session_id}")
+def step(session_id: str, action: Action):
+    env = _get_session(session_id)
+    result = env.step(action)
+    return result.model_dump()
+
+
+@app.get("/state")
+def state_compat(session_id: Optional[str] = Query(default=None)):
+    env = _get_session(_resolve_session_id(session_id))
+    return env.state()
+
+
+@app.get("/state/{session_id}")
+def state(session_id: str):
+    env = _get_session(session_id)
+    return env.state()
+
+
+@app.get("/metadata")
+def metadata():
+    return {
+        "name": "data-cleaning-benchmark",
+        "version": "1.0.0",
+        "description": "LLM agent benchmark for real-world data cleaning tasks.",
+        "tasks": list_task_specs(),
+        "score_range": {
+            "min": DataCleaningEnv.MIN_EPISODE_SCORE,
+            "max": DataCleaningEnv.MAX_EPISODE_SCORE,
+        },
+        "entrypoints": {
+            "reset": "/reset",
+            "step": "/step",
+            "state": "/state",
+            "health": "/health",
+            "tasks": "/tasks",
+            "schema": "/schema",
+            "mcp": "/mcp",
+        },
+    }
+
+
+@app.get("/schema")
+def schema():
+    return {
+        "action": Action.model_json_schema(),
+        "observation": Observation.model_json_schema(),
+        "step_result": StepResult.model_json_schema(),
+        "reset_request": ResetRequest.model_json_schema(),
+    }
+
+
+@app.api_route("/mcp", methods=["GET", "POST"])
+def mcp_metadata():
+    return {
+        "supported": False,
+        "message": "This benchmark exposes simulation HTTP endpoints (reset/step/state).",
+    }
+
+
+@app.delete("/session/{session_id}")
+def delete_session(session_id: str):
+    sessions.pop(session_id, None)
+    return {"deleted": session_id}
+
+
+@app.get("/tasks")
+def list_tasks():
+    return {"tasks": list_task_specs()}
+
+
+def _resolve_session_id(session_id: Optional[str]) -> str:
+    if session_id:
+        return session_id
+    if len(sessions) == 1:
+        return next(iter(sessions.keys()))
+    raise HTTPException(
+        status_code=400,
+        detail="session_id is required when there is not exactly one active session",
+    )
+
+
+def _get_session(session_id: str) -> DataCleaningEnv:
+    env = sessions.get(session_id)
+    if env is None:
+        raise HTTPException(status_code=404, detail=f"Session '{session_id}' not found")
+    return env

env/__init__.py

@@ -0,0 +1 @@
+

env/environment.py

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import numpy as np
+import pandas as pd
+
+from .graders import grade_task
+from .models import Action, Observation, StepResult, TablePreview
+from .rewards import compute_reward
+from .tasks import TASK_IDS, get_task
+
+
+class DataCleaningEnv:
+    MAX_STEPS: int = 20
+    MIN_EPISODE_SCORE: float = 0.01
+    MAX_EPISODE_SCORE: float = 0.99
+
+    def __init__(self) -> None:
+        self.task_id: Optional[str] = None
+        self._task_config: Optional[dict] = None
+        self.original_df: Optional[pd.DataFrame] = None
+        self.current_df: Optional[pd.DataFrame] = None
+        self.step_count: int = 0
+        self.cleaning_log: list = []
+        self.action_history: list = []
+        self.raw_cumulative_reward: float = 0.0
+        self.cumulative_reward: float = 0.0
+        self.done: bool = False
+        self.final_score: float = 0.01
+
+    def reset(self, task_id: Optional[str] = None) -> Observation:
+        if task_id is None:
+            task_id = TASK_IDS[0]
+        self.task_id = task_id
+        self._task_config = get_task(task_id)
+        self.original_df = self._task_config["dirty_df"].copy()
+        self.current_df = self._task_config["dirty_df"].copy()
+        self.step_count = 0
+        self.cleaning_log = []
+        self.action_history = []
+        self.raw_cumulative_reward = 0.0
+        self.cumulative_reward = 0.0
+        self.done = False
+        self.final_score = 0.01
+        return self._build_observation()
+
+    def step(self, action: Action) -> StepResult:
+        if self.done:
+            return StepResult(
+                observation=self._build_observation(),
+                reward=self.final_score,
+                done=True,
+                info={
+                    "error": "Episode already finished",
+                    "cumulative_reward": self.cumulative_reward,
+                    "raw_cumulative_reward": self.raw_cumulative_reward,
+                    "final_score": self.final_score,
+                    "step": self.step_count,
+                },
+            )
+
+        error: Optional[str] = None
+        reward: float = 0.0
+
+        if action.type == "submit":
+            self.final_score = grade_task(self.task_id, self.current_df)
+            reward = self.final_score
+            self.cleaning_log.append(f"[SUBMIT] Final grade: {self.final_score:.4f}")
+            self.done = True
+        else:
+            try:
+                reward, log_msg = self._apply_action(action)
+                self.cleaning_log.append(log_msg)
+            except Exception as exc:
+                error = str(exc)
+                reward = -0.10
+                self.cleaning_log.append(f"[ERROR] {error}")
+
+        self.step_count += 1
+        self.raw_cumulative_reward = round(self.raw_cumulative_reward + reward, 4)
+        self.cumulative_reward = self._clamp_episode_score(self.raw_cumulative_reward)
+        self.action_history.append(action.model_dump())
+
+        if not self.done and self.step_count >= self.MAX_STEPS:
+            self.final_score = grade_task(self.task_id, self.current_df)
+            self.done = True
+
+        return StepResult(
+            observation=self._build_observation(),
+            reward=round(reward, 4),
+            done=self.done,
+            info={
+                "error": error,
+                "cumulative_reward": self.cumulative_reward,
+                "raw_cumulative_reward": self.raw_cumulative_reward,
+                "final_score": self.final_score,
+                "step": self.step_count,
+            },
+        )
+
+    def state(self) -> dict:
+        return {
+            "task_id": self.task_id,
+            "step_count": self.step_count,
+            "cumulative_reward": self.cumulative_reward,
+            "raw_cumulative_reward": self.raw_cumulative_reward,
+            "final_score": self.final_score,
+            "done": self.done,
+            "cleaning_log": self.cleaning_log,
+            "action_history": self.action_history,
+            "current_data": self._df_records_with_none(self.current_df) if self.current_df is not None else [],
+        }
+
+    @classmethod
+    def _clamp_episode_score(cls, value: float) -> float:
+        return round(min(max(value, cls.MIN_EPISODE_SCORE), cls.MAX_EPISODE_SCORE), 4)
+
+    def _apply_action(self, action: Action) -> Tuple[float, str]:
+        df = self.current_df
+
+        if action.type == "fill_missing":
+            col = self._require_column(action.column, df)
+            missing_before = int(df[col].isna().sum())
+            if missing_before == 0:
+                return -0.05, f"[WARN] No missing values in '{col}' — wasted step"
+
+            if action.strategy == "mean":
+                df[col] = df[col].fillna(df[col].mean())
+            elif action.strategy == "median":
+                df[col] = df[col].fillna(df[col].median())
+            elif action.strategy == "mode":
+                df[col] = df[col].fillna(df[col].mode().iloc[0])
+            elif action.strategy == "constant":
+                df[col] = df[col].fillna(action.value)
+            else:
+                raise ValueError(f"Unknown fill strategy '{action.strategy}'")
+
+            reward = compute_reward("fill_missing", {"filled": missing_before})
+            return reward, f"Filled {missing_before} missing values in '{col}' via {action.strategy}"
+
+        if action.type == "standardize_values":
+            col = self._require_column(action.column, df)
+            if not action.mapping:
+                raise ValueError("'mapping' dict is required for standardize_values")
+            replaced = int(df[col].isin(action.mapping.keys()).sum())
+            df[col] = df[col].apply(lambda x: action.mapping.get(str(x), x) if pd.notna(x) else x)
+            reward = compute_reward("standardize_values", {"replaced": replaced})
+            return reward, f"Standardised {replaced} values in '{col}'"
+
+        if action.type == "remove_duplicates":
+            before = len(df)
+            self.current_df = df.drop_duplicates().reset_index(drop=True)
+            removed = before - len(self.current_df)
+            if removed == 0:
+                return -0.05, "[WARN] No exact duplicates found — wasted step"
+            reward = compute_reward("remove_duplicates", {"removed": removed})
+            return reward, f"Removed {removed} duplicate row(s)"
+
+        if action.type == "remove_row":
+            if action.row_id is None:
+                raise ValueError("'row_id' is required for remove_row")
+            if action.row_id not in df.index:
+                raise ValueError(f"Row index {action.row_id} not found (valid range 0–{len(df)-1})")
+            self.current_df = df.drop(index=action.row_id).reset_index(drop=True)
+            reward = compute_reward("remove_row", {})
+            return reward, f"Removed row at index {action.row_id}"
+
+        if action.type == "convert_type":
+            col = self._require_column(action.column, df)
+            tgt = action.target_type
+
+            if tgt == "float":
+                df[col] = (
+                    df[col]
+                    .astype(str)
+                    .str.replace(r"[$,\s]", "", regex=True)
+                    .replace("nan", np.nan)
+                    .replace("None", np.nan)
+                )
+                df[col] = pd.to_numeric(df[col], errors="coerce")
+            elif tgt == "int":
+                df[col] = pd.to_numeric(df[col], errors="coerce").astype("Int64")
+            elif tgt == "str":
+                df[col] = df[col].astype(str)
+            elif tgt == "datetime":
+                parsed = pd.to_datetime(df[col], errors="coerce")
+                df[col] = parsed.dt.strftime("%Y-%m-%d")
+            else:
+                raise ValueError(f"Unknown target_type '{tgt}'")
+
+            reward = compute_reward("convert_type", {})
+            return reward, f"Converted column '{col}' → {tgt}"
+
+        if action.type == "clip_outliers":
+            col = self._require_column(action.column, df)
+            if action.lower is None and action.upper is None:
+                raise ValueError("At least one of 'lower' or 'upper' must be set")
+
+            series = pd.to_numeric(df[col], errors="coerce")
+            clipped = 0
+            if action.lower is not None:
+                clipped += int((series < action.lower).sum())
+            if action.upper is not None:
+                clipped += int((series > action.upper).sum())
+
+            df[col] = series.clip(lower=action.lower, upper=action.upper)
+            reward = compute_reward("clip_outliers", {"clipped": clipped})
+            return reward, f"Clipped '{col}' to [{action.lower}, {action.upper}] ({clipped} value(s) affected)"
+
+        raise ValueError(f"Unknown action type '{action.type}'")
+
+    @staticmethod
+    def _require_column(col: Optional[str], df: pd.DataFrame) -> str:
+        if not col:
+            raise ValueError("'column' field is required for this action")
+        if col not in df.columns:
+            raise ValueError(f"Column '{col}' not found. Available: {list(df.columns)}")
+        return col
+
+    @staticmethod
+    def _df_records_with_none(df: pd.DataFrame) -> list[dict]:
+        safe_df = df.astype(object).where(pd.notna(df), None)
+        return safe_df.to_dict(orient="records")
+
+    def _build_observation(self) -> Observation:
+        df = self.current_df
+        issues: list = []
+
+        if df is not None:
+            for col in df.columns:
+                miss = int(df[col].isna().sum())
+                if miss > 0:
+                    issues.append(f"Column '{col}' has {miss} missing value(s)")
+            dup = int(df.duplicated().sum())
+            if dup > 0:
+                issues.append(f"{dup} exact duplicate row(s) detected")
+
+            head = df.head(10).copy()
+            head.insert(0, "_row_id", head.index.tolist())
+            preview_rows = self._df_records_with_none(head)
+            schema_info = {c: str(df[c].dtype) for c in df.columns}
+            shape = list(df.shape)
+        else:
+            preview_rows, schema_info, shape = [], {}, [0, 0]
+
+        preview = TablePreview(
+            columns=["_row_id"] + (list(df.columns) if df is not None else []),
+            rows=preview_rows,
+            shape=shape,
+        )
+
+        return Observation(
+            task_id=self.task_id or "",
+            task_description=(self._task_config["description"] if self._task_config else ""),
+            table_preview=preview,
+            schema_info=schema_info,
+            valid_actions=[
+                "fill_missing",
+                "standardize_values",
+                "remove_duplicates",
+                "remove_row",
+                "convert_type",
+                "clip_outliers",
+                "submit",
+            ],
+            step=self.step_count,
+            max_steps=self.MAX_STEPS,
+            cleaning_log=self.cleaning_log[-6:],
+            issues_detected=issues,
+        )

env/graders.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import math
+
+import pandas as pd
+
+
+def _strict_score(value: float) -> float:
+    try:
+        score = float(value)
+    except (TypeError, ValueError):
+        return 0.01
+
+    if not math.isfinite(score):
+        return 0.01
+
+    return round(min(max(score, 0.01), 0.99), 4)
+
+
+def grade_task1(df: pd.DataFrame) -> float:
+    score = 0.0
+
+    if df.duplicated().sum() == 0:
+        score += 0.25
+
+    if "email" in df.columns and df["email"].isna().sum() == 0:
+        score += 0.25
+
+    if "age" in df.columns and df["age"].isna().sum() == 0:
+        score += 0.25
+
+    valid_countries = {"United States", "United Kingdom", "Canada", "Australia"}
+    if "country" in df.columns:
+        non_null = df["country"].dropna()
+        if len(non_null) == 0:
+            pass
+        elif set(non_null.unique()).issubset(valid_countries):
+            score += 0.25
+        else:
+            valid_n = non_null.isin(valid_countries).sum()
+            score += 0.25 * (valid_n / len(non_null))
+
+    return _strict_score(score)
+
+
+def grade_task2(df: pd.DataFrame) -> float:
+    score = 0.0
+    n = len(df)
+    if n == 0:
+        return 0.01
+
+    if "date" in df.columns:
+        pattern = r"^\d{4}-\d{2}-\d{2}$"
+        valid = df["date"].astype(str).str.match(pattern).sum()
+        score += 0.25 * (valid / n)
+
+    if "price" in df.columns:
+        numeric = pd.to_numeric(df["price"], errors="coerce")
+        non_null = numeric.notna().sum()
+        score += 0.25 * (non_null / n)
+
+    valid_cats = {"Electronics", "Furniture"}
+    if "category" in df.columns:
+        non_null_cats = df["category"].dropna()
+        if len(non_null_cats) > 0:
+            valid_n = non_null_cats.isin(valid_cats).sum()
+            score += 0.25 * (valid_n / len(non_null_cats))
+
+    key_cols = [c for c in ["price", "category", "quantity"] if c in df.columns]
+    if key_cols:
+        total_cells = n * len(key_cols)
+        missing = sum(int(df[c].isna().sum()) for c in key_cols)
+        score += 0.25 * (1.0 - missing / total_cells)
+
+    return _strict_score(score)
+
+
+def grade_task3(df: pd.DataFrame) -> float:
+    score = 0.0
+    n = len(df)
+    if n == 0:
+        return 0.01
+
+    if "user_id" in df.columns:
+        dup = df["user_id"].duplicated().sum()
+        if dup == 0:
+            score += 0.34
+        else:
+            score += 0.34 * (1.0 - dup / n)
+
+    if "session_duration" in df.columns:
+        max_dur = df["session_duration"].dropna().max() if n > 0 else 0
+        if max_dur <= 1000:
+            score += 0.33
+        elif max_dur <= 5000:
+            score += 0.15
+
+    if "bounce_rate" in df.columns:
+        valid_br = ((df["bounce_rate"] >= 0) & (df["bounce_rate"] <= 1)).sum()
+        score += 0.165 * (valid_br / n)
+
+    if "page_views" in df.columns and df["page_views"].isna().sum() == 0:
+        score += 0.165
+
+    return _strict_score(score)
+
+
+def grade_task(task_id: str, df: pd.DataFrame) -> float:
+    fn = TASK_GRADERS.get(task_id)
+    if fn is None:
+        return 0.01
+    return fn(df)
+
+
+def grade_task1_easy(df: pd.DataFrame) -> float:
+    return grade_task1(df)
+
+
+def grade_task2_medium(df: pd.DataFrame) -> float:
+    return grade_task2(df)
+
+
+def grade_task3_hard(df: pd.DataFrame) -> float:
+    return grade_task3(df)
+
+
+def grade_task4_medium_alt(df: pd.DataFrame) -> float:
+    return grade_task2(df)
+
+
+def grade_task5_hard_alt(df: pd.DataFrame) -> float:
+    return grade_task3(df)
+
+
+TASK_GRADERS = {
+    "task1_easy": grade_task1_easy,
+    "task2_medium": grade_task2_medium,
+    "task3_hard": grade_task3_hard,
+    "task4_medium_alt": grade_task4_medium_alt,
+    "task5_hard_alt": grade_task5_hard_alt,
+}

env/models.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel
+
+
+class Action(BaseModel):
+    type: str
+    column: Optional[str] = None
+    row_id: Optional[int] = None
+    strategy: Optional[str] = None
+    value: Optional[Union[str, float, int]] = None
+    mapping: Optional[Dict[str, str]] = None
+    target_type: Optional[str] = None
+    lower: Optional[float] = None
+    upper: Optional[float] = None
+
+
+class TablePreview(BaseModel):
+    columns: List[str]
+    rows: List[Dict[str, Any]]
+    shape: List[int]
+
+
+class Observation(BaseModel):
+    task_id: str
+    task_description: str
+    table_preview: TablePreview
+    schema_info: Dict[str, str]
+    valid_actions: List[str]
+    step: int
+    max_steps: int
+    cleaning_log: List[str]
+    issues_detected: List[str]
+
+
+class StepResult(BaseModel):
+    observation: Observation
+    reward: float
+    done: bool
+    info: Dict[str, Any]

env/rewards.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+
+def compute_reward(action_type: str, context: dict) -> float:
+    """
+    Intermediate reward shaping.
+    Final episode reward comes from the grader (called at submit).
+    """
+    if action_type == "fill_missing":
+        filled = context.get("filled", 0)
+        return round(min(0.08 * filled, 0.30), 4)
+
+    if action_type == "standardize_values":
+        replaced = context.get("replaced", 0)
+        return round(min(0.06 * replaced, 0.25), 4)
+
+    if action_type == "remove_duplicates":
+        removed = context.get("removed", 0)
+        return round(min(0.15 * removed, 0.30), 4)
+
+    if action_type == "remove_row":
+        return 0.05
+
+    if action_type == "convert_type":
+        return 0.15
+
+    if action_type == "clip_outliers":
+        clipped = context.get("clipped", 0)
+        return round(min(0.10 * max(clipped, 1), 0.30), 4)
+
+    return 0.0

env/tasks.py

@@ -0,0 +1,291 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+import pandas as pd
+
+
+TASK1_DIRTY = [
+    {"name": "Alice Johnson", "email": "alice@email.com", "country": "USA", "age": 28.0},
+    {"name": "Bob Smith", "email": "bob@email.com", "country": "United States", "age": None},
+    {"name": "Carol White", "email": "carol@email.com", "country": "UK", "age": 35.0},
+    {"name": "Alice Johnson", "email": "alice@email.com", "country": "USA", "age": 28.0},
+    {"name": "Dave Brown", "email": None, "country": "US", "age": 42.0},
+    {"name": "Eve Davis", "email": "eve@email.com", "country": "United Kingdom", "age": 31.0},
+    {"name": "Frank Miller", "email": "frank@email.com", "country": "Canada", "age": None},
+    {"name": "Grace Wilson", "email": "grace@email.com", "country": "CAN", "age": 25.0},
+    {"name": "Henry Moore", "email": "henry@email.com", "country": "australia", "age": 38.0},
+    {"name": "Iris Taylor", "email": "iris@email.com", "country": "AUS", "age": 29.0},
+]
+
+TASK1_DESCRIPTION = (
+    "Clean a customer dataset. Issues to fix:\n"
+    "1) Remove exact duplicate rows\n"
+    "2) Fill missing emails using constant 'unknown@email.com'\n"
+    "3) Fill missing ages using median\n"
+    "4) Standardize country names to United States, United Kingdom, Canada, Australia"
+)
+
+TASK2_DIRTY = [
+    {
+        "order_id": 1,
+        "date": "2023-01-15",
+        "product": "Laptop",
+        "category": "Electronics",
+        "price": "$1200.00",
+        "quantity": 2,
+    },
+    {
+        "order_id": 2,
+        "date": "02/20/2023",
+        "product": "Chair",
+        "category": "Furniture",
+        "price": "$250.50",
+        "quantity": 1,
+    },
+    {
+        "order_id": 3,
+        "date": "Mar 10, 2023",
+        "product": "Headphones",
+        "category": "Electronics",
+        "price": "$89.99",
+        "quantity": 3,
+    },
+    {
+        "order_id": 4,
+        "date": "2023-04-05",
+        "product": "Desk",
+        "category": "Furnitre",
+        "price": "$450.00",
+        "quantity": 1,
+    },
+    {
+        "order_id": 5,
+        "date": "05/12/2023",
+        "product": "Monitor",
+        "category": "Electronics",
+        "price": "320.00",
+        "quantity": 2,
+    },
+    {
+        "order_id": 6,
+        "date": "2023-06-18",
+        "product": "Keyboard",
+        "category": None,
+        "price": "$75.00",
+        "quantity": 5,
+    },
+    {
+        "order_id": 7,
+        "date": "July 22 2023",
+        "product": "Mouse",
+        "category": "Electronics",
+        "price": "$35.00",
+        "quantity": 4,
+    },
+    {
+        "order_id": 8,
+        "date": "2023-08-30",
+        "product": "Bookshelf",
+        "category": "Furniture",
+        "price": None,
+        "quantity": 1,
+    },
+    {
+        "order_id": 9,
+        "date": "09-14-2023",
+        "product": "Webcam",
+        "category": "ELECTRONICS",
+        "price": "$65.00",
+        "quantity": 2,
+    },
+    {
+        "order_id": 10,
+        "date": "2023-10-01",
+        "product": "Lamp",
+        "category": "Furniture",
+        "price": "$45.00",
+        "quantity": 3,
+    },
+    {
+        "order_id": 11,
+        "date": "11/15/2023",
+        "product": "Tablet",
+        "category": "Electronix",
+        "price": "$599.00",
+        "quantity": 1,
+    },
+    {
+        "order_id": 12,
+        "date": "2023-12-20",
+        "product": "Sofa",
+        "category": "Furniture",
+        "price": "$1100.00",
+        "quantity": 1,
+    },
+]
+
+TASK2_DESCRIPTION = (
+    "Clean an e-commerce orders dataset. Issues to fix:\n"
+    "1) Normalise all dates to YYYY-MM-DD format using convert_type(date, datetime)\n"
+    "2) Convert price column to float (strips $ signs automatically)\n"
+    "3) Standardise category typos: Furnitre to Furniture, ELECTRONICS to Electronics, Electronix to Electronics\n"
+    "4) Fill missing price with median; fill or remove missing category rows"
+)
+
+TASK3_DIRTY = [
+    {"user_id": "U001", "name": "Alice Johnson", "page_views": 45, "session_duration": 320, "bounce_rate": 0.25},
+    {"user_id": "U001", "name": "Alice J.", "page_views": 45, "session_duration": 315, "bounce_rate": 0.25},
+    {"user_id": "U002", "name": "Bob Smith", "page_views": 12, "session_duration": 85000, "bounce_rate": 0.80},
+    {"user_id": "U003", "name": "Carol White", "page_views": 67, "session_duration": 450, "bounce_rate": 0.15},
+    {"user_id": "U004", "name": "Dave Brown", "page_views": 23, "session_duration": 190, "bounce_rate": 0.55},
+    {"user_id": "U005", "name": "Eve Davis", "page_views": 89, "session_duration": 95000, "bounce_rate": 0.10},
+    {"user_id": "U003", "name": "Carol White", "page_views": 67, "session_duration": 450, "bounce_rate": 0.15},
+    {"user_id": "U006", "name": "Frank Miller", "page_views": None, "session_duration": 280, "bounce_rate": 0.45},
+    {"user_id": "U007", "name": "Grace Wilson", "page_views": 34, "session_duration": 360, "bounce_rate": 1.50},
+    {"user_id": "U008", "name": "Henry Moore", "page_views": 56, "session_duration": 420, "bounce_rate": 0.35},
+    {"user_id": "U009", "name": "Iris Taylor", "page_views": 78, "session_duration": 78000, "bounce_rate": 0.20},
+    {"user_id": "U010", "name": "Jack Wilson", "page_views": 19, "session_duration": 150, "bounce_rate": 0.70},
+]
+
+TASK3_DESCRIPTION = (
+    "Clean a web analytics dataset. Issues to fix:\n"
+    "1) Remove duplicate user_ids (exact + near-duplicates, keep first occurrence)\n"
+    "2) Clip session_duration outliers to max 1000 seconds\n"
+    "3) Clip bounce_rate to valid range [0.0, 1.0]\n"
+    "4) Fill missing page_views with median"
+)
+
+TASK4_DESCRIPTION = (
+    "Alternative medium data-cleaning scenario based on e-commerce orders.\n"
+    "Use the same cleaning operations as task2_medium and submit a clean table."
+)
+
+TASK5_DESCRIPTION = (
+    "Alternative hard data-cleaning scenario based on analytics logs.\n"
+    "Use the same cleaning operations as task3_hard and submit a clean table."
+)
+
+TASK_GRADER_ENTRYPOINTS_COLON = {
+    "task1_easy": "env.graders:grade_task1_easy",
+    "task2_medium": "env.graders:grade_task2_medium",
+    "task3_hard": "env.graders:grade_task3_hard",
+    "task4_medium_alt": "env.graders:grade_task4_medium_alt",
+    "task5_hard_alt": "env.graders:grade_task5_hard_alt",
+}
+
+TASK_GRADER_ENTRYPOINTS_DOTTED = {
+    "task1_easy": "env.graders.grade_task1_easy",
+    "task2_medium": "env.graders.grade_task2_medium",
+    "task3_hard": "env.graders.grade_task3_hard",
+    "task4_medium_alt": "env.graders.grade_task4_medium_alt",
+    "task5_hard_alt": "env.graders.grade_task5_hard_alt",
+}
+
+
+def get_task(task_id: str) -> Dict[str, Any]:
+    registry = {
+        "task1_easy": {
+            "description": TASK1_DESCRIPTION,
+            "dirty_df": pd.DataFrame(TASK1_DIRTY),
+            "task_id": "task1_easy",
+            "difficulty": "easy",
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task1_easy"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task1_easy"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task1_easy"],
+        },
+        "task2_medium": {
+            "description": TASK2_DESCRIPTION,
+            "dirty_df": pd.DataFrame(TASK2_DIRTY),
+            "task_id": "task2_medium",
+            "difficulty": "medium",
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task2_medium"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task2_medium"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task2_medium"],
+        },
+        "task3_hard": {
+            "description": TASK3_DESCRIPTION,
+            "dirty_df": pd.DataFrame(TASK3_DIRTY),
+            "task_id": "task3_hard",
+            "difficulty": "hard",
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task3_hard"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task3_hard"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task3_hard"],
+        },
+        "task4_medium_alt": {
+            "description": TASK4_DESCRIPTION,
+            "dirty_df": pd.DataFrame(TASK2_DIRTY),
+            "task_id": "task4_medium_alt",
+            "difficulty": "medium",
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task4_medium_alt"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task4_medium_alt"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task4_medium_alt"],
+        },
+        "task5_hard_alt": {
+            "description": TASK5_DESCRIPTION,
+            "dirty_df": pd.DataFrame(TASK3_DIRTY),
+            "task_id": "task5_hard_alt",
+            "difficulty": "hard",
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task5_hard_alt"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task5_hard_alt"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task5_hard_alt"],
+        },
+    }
+    if task_id not in registry:
+        raise ValueError(f"Unknown task_id '{task_id}'. Choose from: {list(registry)}")
+    cfg = registry[task_id]
+    cfg["dirty_df"] = cfg["dirty_df"].copy()
+    return cfg
+
+
+TASK_IDS = ["task1_easy", "task2_medium", "task3_hard", "task4_medium_alt", "task5_hard_alt"]
+
+
+def list_tasks() -> list[dict[str, Any]]:
+    return [
+        {
+            "id": "task1_easy",
+            "task_id": "task1_easy",
+            "difficulty": "easy",
+            "max_steps": 20,
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task1_easy"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task1_easy"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task1_easy"],
+        },
+        {
+            "id": "task2_medium",
+            "task_id": "task2_medium",
+            "difficulty": "medium",
+            "max_steps": 20,
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task2_medium"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task2_medium"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task2_medium"],
+        },
+        {
+            "id": "task3_hard",
+            "task_id": "task3_hard",
+            "difficulty": "hard",
+            "max_steps": 20,
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task3_hard"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task3_hard"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task3_hard"],
+        },
+        {
+            "id": "task4_medium_alt",
+            "task_id": "task4_medium_alt",
+            "difficulty": "medium",
+            "max_steps": 20,
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task4_medium_alt"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task4_medium_alt"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task4_medium_alt"],
+        },
+        {
+            "id": "task5_hard_alt",
+            "task_id": "task5_hard_alt",
+            "difficulty": "hard",
+            "max_steps": 20,
+            "grader": TASK_GRADER_ENTRYPOINTS_DOTTED["task5_hard_alt"],
+            "grader_fn": TASK_GRADER_ENTRYPOINTS_COLON["task5_hard_alt"],
+            "grader_path": TASK_GRADER_ENTRYPOINTS_COLON["task5_hard_alt"],
+        },
+    ]

inference.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+
+from dotenv import load_dotenv
+from openai import OpenAI
+
+from env.environment import DataCleaningEnv
+from env.models import Action
+
+load_dotenv()
+
+API_BASE_URL = os.getenv("API_BASE_URL", "https://api.groq.com/openai/v1")
+MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/llama-4-scout-17b-16e-instruct")
+HF_TOKEN = os.getenv("HF_TOKEN")
+TASK_ID = os.getenv("TASK_ID", "task1_easy")
+MAX_STEPS = int(os.getenv("MAX_STEPS", "15"))
+ENV_NAME = "data-cleaning-benchmark"
+
+if HF_TOKEN is None:
+    raise ValueError("HF_TOKEN environment variable is required")
+
+client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
+
+SYSTEM_PROMPT = """You are a data cleaning agent. Analyse the observation and choose ONE cleaning action.
+
+Available action types and required fields:
+  fill_missing       -> column (str), strategy (mean|median|mode|constant), value (if constant)
+  standardize_values -> column (str), mapping (dict old->new)
+  remove_duplicates  -> (no extra fields)
+  remove_row         -> row_id (int from _row_id column in preview)
+  convert_type       -> column (str), target_type (float|int|str|datetime)
+  clip_outliers      -> column (str), lower (float|null), upper (float|null)
+  submit             -> (no extra fields; use when dataset is clean)
+
+Rules:
+- Respond with a SINGLE valid JSON object and NOTHING else.
+- No markdown fences, no explanation.
+- When no issues remain, always respond with: {"type": "submit"}
+
+Examples:
+{"type": "remove_duplicates"}
+{"type": "fill_missing", "column": "age", "strategy": "median"}
+{"type": "standardize_values", "column": "country", "mapping": {"USA": "United States", "US": "United States", "UK": "United Kingdom", "CAN": "Canada", "australia": "Australia", "AUS": "Australia"}}
+{"type": "convert_type", "column": "date", "target_type": "datetime"}
+{"type": "convert_type", "column": "price", "target_type": "float"}
+{"type": "clip_outliers", "column": "session_duration", "lower": 0.0, "upper": 1000.0}
+{"type": "submit"}
+"""
+
+
+def get_action(obs_dict: dict, history: list[dict]) -> dict:
+    user_msg = {
+        "role": "user",
+        "content": (
+            "Current observation:\n" + json.dumps(obs_dict, indent=2, default=str) + "\n\nNext action (JSON only):"
+        ),
+    }
+    history.append(user_msg)
+
+    response = client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=[{"role": "system", "content": SYSTEM_PROMPT}] + history,
+        max_tokens=256,
+        temperature=0,
+    )
+    raw = response.choices[0].message.content.strip()
+    history.append({"role": "assistant", "content": raw})
+
+    clean = re.sub(r"```[a-z]*\n?", "", raw).replace("```", "").strip()
+    try:
+        return json.loads(clean)
+    except json.JSONDecodeError:
+        match = re.search(r"\{.*\}", clean, re.DOTALL)
+        if match:
+            return json.loads(match.group())
+        return {"type": "submit"}
+
+
+def run_inference() -> None:
+    env = DataCleaningEnv()
+    rewards: list[float] = []
+    history: list[dict] = []
+    step = 0
+    done = False
+    success = False
+
+    print(f"[START] task={TASK_ID} env={ENV_NAME} model={MODEL_NAME}", flush=True)
+
+    try:
+        obs = env.reset(task_id=TASK_ID)
+
+        while not done and step < MAX_STEPS:
+            try:
+                action_dict = get_action(obs.model_dump(), history)
+                action = Action(**action_dict)
+            except Exception:
+                action_dict = {"type": "submit"}
+                action = Action(type="submit")
+
+            result = env.step(action)
+            obs = result.observation
+            done = result.done
+            reward = result.reward
+            error = result.info.get("error")
+
+            rewards.append(reward)
+            step += 1
+
+            action_str = json.dumps(action_dict, separators=(",", ":"), default=str)
+            print(
+                f"[STEP] step={step} action={action_str} "
+                f"reward={reward:.2f} done={'true' if done else 'false'} "
+                f"error={error if error else 'null'}",
+                flush=True,
+            )
+
+        if not done:
+            result = env.step(Action(type="submit"))
+            rewards.append(result.reward)
+            step += 1
+            print(
+                f"[STEP] step={step} action={{\"type\":\"submit\"}} "
+                f"reward={result.reward:.2f} done=true error={result.info.get('error') or 'null'}",
+                flush=True,
+            )
+
+        success = bool(env.final_score >= 0.5)
+    except Exception:
+        success = False
+    finally:
+        try:
+            if hasattr(env, "close"):
+                env.close()
+        except Exception:
+            pass
+
+        rewards_str = ",".join(f"{reward:.2f}" for reward in rewards)
+        print(
+            f"[END] success={'true' if success else 'false'} "
+            f"steps={step} score={env.final_score:.2f} rewards={rewards_str}",
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    run_inference()

openenv.yaml

@@ -0,0 +1,117 @@
+name: data-cleaning-benchmark
+version: "1.0.0"
+description: >
+  A multi-task LLM agent benchmark for real-world tabular data cleaning.
+  The agent receives a dirty dataset and must apply structured cleaning
+  actions to fix duplicates, missing values, format issues, and outliers.
+
+author: "Jayesh"
+license: MIT
+
+tasks:
+  - id: task1_easy
+    task_id: task1_easy
+    name: "Basic Customer Data Cleanup"
+    difficulty: easy
+    max_steps: 20
+    description: "Remove duplicates, fill missing values, standardise country names."
+    grader: env.graders.grade_task1_easy
+    grader_fn: env.graders:grade_task1_easy
+    grader_path: env.graders:grade_task1_easy
+  - id: task2_medium
+    task_id: task2_medium
+    name: "E-commerce Orders Normalisation"
+    difficulty: medium
+    max_steps: 20
+    description: "Fix mixed date formats, convert price strings, correct category typos."
+    grader: env.graders.grade_task2_medium
+    grader_fn: env.graders:grade_task2_medium
+    grader_path: env.graders:grade_task2_medium
+  - id: task3_hard
+    task_id: task3_hard
+    name: "Analytics Data Deep Clean"
+    difficulty: hard
+    max_steps: 20
+    description: "Resolve duplicate user IDs, clip session outliers, fix invalid bounce rates."
+    grader: env.graders.grade_task3_hard
+    grader_fn: env.graders:grade_task3_hard
+    grader_path: env.graders:grade_task3_hard
+  - id: task4_medium_alt
+    task_id: task4_medium_alt
+    name: "E-commerce Orders Cleanup (Alt)"
+    difficulty: medium
+    max_steps: 20
+    description: "Alternative medium scenario sharing the same grading criteria as task2_medium."
+    grader: env.graders.grade_task4_medium_alt
+    grader_fn: env.graders:grade_task4_medium_alt
+    grader_path: env.graders:grade_task4_medium_alt
+  - id: task5_hard_alt
+    task_id: task5_hard_alt
+    name: "Analytics Deep Clean (Alt)"
+    difficulty: hard
+    max_steps: 20
+    description: "Alternative hard scenario sharing the same grading criteria as task3_hard."
+    grader: env.graders.grade_task5_hard_alt
+    grader_fn: env.graders:grade_task5_hard_alt
+    grader_path: env.graders:grade_task5_hard_alt
+
+observation_space:
+  type: structured_json
+  fields:
+    - task_id
+    - task_description
+    - table_preview
+    - schema_info
+    - valid_actions
+    - step / max_steps
+    - cleaning_log
+    - issues_detected
+
+action_space:
+  type: structured_json
+  actions:
+    - name: fill_missing
+      params: ["column", "strategy(mean|median|mode|constant)", "value?"]
+    - name: standardize_values
+      params: ["column", "mapping(dict)"]
+    - name: remove_duplicates
+      params: []
+    - name: remove_row
+      params: ["row_id(int)"]
+    - name: convert_type
+      params: ["column", "target_type(float|int|str|datetime)"]
+    - name: clip_outliers
+      params: ["column", "lower?", "upper?"]
+    - name: submit
+      params: []
+
+reward:
+  type: shaped
+  intermediate: true
+  range: [0.01, 0.99]
+  description: >
+    Positive rewards for correct cleaning steps; small penalties for
+    invalid or wasted actions; final grader score awarded on submit().
+
+api:
+  base_path: "/"
+  endpoints:
+    reset: "POST /reset"
+    step: "POST /step"
+    state: "GET  /state"
+    step_legacy: "POST /step/{session_id}"
+    state_legacy: "GET  /state/{session_id}"
+    health: "GET  /health"
+    tasks: "GET  /tasks"
+
+runtime:
+  language: python
+  version: "3.11"
+  port: 7860
+  framework: fastapi
+
+tags:
+  - openenv
+  - data-cleaning
+  - llm-benchmark
+  - tabular

pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "data-cleaning-benchmark"
+version = "1.0.0"
+description = "A multi-task OpenEnv benchmark for tabular data cleaning."
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "fastapi==0.110.0",
+    "uvicorn==0.27.1",
+    "pydantic==2.6.3",
+    "pandas==2.2.1",
+    "numpy==1.26.4",
+    "openai>=2.7.2",
+    "openenv>=0.2.0",
+    "python-dotenv==1.0.1",
+    "httpx==0.27.0",
+    "pytest==8.1.1",
+]
+
+[project.scripts]
+server = "server.cli:main"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

requirements.txt

@@ -0,0 +1,10 @@
+fastapi==0.110.0
+uvicorn==0.27.1
+pydantic==2.6.3
+pandas==2.2.1
+numpy==1.26.4
+openai>=2.7.2
+openenv>=0.2.0
+python-dotenv==1.0.1
+httpx==0.27.0
+pytest==8.1.1

server/__init__.py

@@ -0,0 +1 @@
+

server/app.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import os
+
+import uvicorn
+
+from app import app
+
+
+def main() -> None:
+    host = os.getenv("HOST", "0.0.0.0")
+    port = int(os.getenv("PORT", "7860"))
+    uvicorn.run("server.app:app", host=host, port=port)
+
+
+if __name__ == "__main__":
+    main()

server/cli.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import os
+
+import uvicorn
+
+
+def main() -> None:
+    host = os.getenv("HOST", "0.0.0.0")
+    port = int(os.getenv("PORT", "7860"))
+    uvicorn.run("server.app:app", host=host, port=port)

tests/test_graders.py

@@ -0,0 +1,70 @@
+import pandas as pd
+
+from env.graders import grade_task, grade_task1, grade_task2, grade_task3
+from env.tasks import get_task
+
+
+def test_grade_task1_dirty_is_low():
+    cfg = get_task("task1_easy")
+    score = grade_task1(cfg["dirty_df"])
+    assert 0.0 < score <= 0.5
+
+
+def test_grade_task1_perfect_is_bounded():
+    df = pd.DataFrame(
+        {
+            "name": ["Alice", "Bob", "Carol"],
+            "email": ["a@x.com", "b@x.com", "c@x.com"],
+            "country": ["United States", "United Kingdom", "Australia"],
+            "age": [28.0, 35.0, 42.0],
+        }
+    )
+    score = grade_task1(df)
+    assert 0.99 == score
+
+
+def test_grade_task1_partial():
+    df = pd.DataFrame(
+        {
+            "name": ["Alice", "Bob"],
+            "email": ["a@x.com", "b@x.com"],
+            "country": ["USA", "UK"],
+            "age": [28.0, 35.0],
+        }
+    )
+    score = grade_task1(df)
+    assert 0.4 < score < 0.99
+
+
+def test_grade_task2_score_range():
+    cfg = get_task("task2_medium")
+    score = grade_task2(cfg["dirty_df"])
+    assert 0.0 < score < 1.0
+
+
+def test_grade_task3_score_range():
+    cfg = get_task("task3_hard")
+    score = grade_task3(cfg["dirty_df"])
+    assert 0.0 < score < 1.0
+
+
+def test_grade_task_dispatcher():
+    for tid in ["task1_easy", "task2_medium", "task3_hard", "task4_medium_alt", "task5_hard_alt"]:
+        cfg = get_task(tid)
+        s = grade_task(tid, cfg["dirty_df"])
+        assert 0.0 < s < 1.0
+
+
+def test_grader_not_constant():
+    cfg = get_task("task1_easy")
+    dirty_score = grade_task1(cfg["dirty_df"])
+    clean_df = pd.DataFrame(
+        {
+            "name": ["Alice", "Bob"],
+            "email": ["a@x.com", "b@x.com"],
+            "country": ["United States", "Australia"],
+            "age": [28.0, 35.0],
+        }
+    )
+    clean_score = grade_task1(clean_df)
+    assert clean_score != dirty_score

tests/test_reset.py

@@ -0,0 +1,50 @@
+import pytest
+
+from env.environment import DataCleaningEnv
+
+
+def test_reset_default():
+    env = DataCleaningEnv()
+    obs = env.reset()
+    assert obs.task_id == "task1_easy"
+    assert obs.step == 0
+    assert obs.max_steps == 20
+    assert len(obs.table_preview.rows) > 0
+    assert "remove_duplicates" in obs.valid_actions
+
+
+def test_reset_task2():
+    env = DataCleaningEnv()
+    obs = env.reset(task_id="task2_medium")
+    assert obs.task_id == "task2_medium"
+    assert obs.step == 0
+
+
+def test_reset_task3():
+    env = DataCleaningEnv()
+    obs = env.reset(task_id="task3_hard")
+    assert obs.task_id == "task3_hard"
+
+
+def test_reset_task4_alt():
+    env = DataCleaningEnv()
+    obs = env.reset(task_id="task4_medium_alt")
+    assert obs.task_id == "task4_medium_alt"
+
+
+def test_reset_task5_alt():
+    env = DataCleaningEnv()
+    obs = env.reset(task_id="task5_hard_alt")
+    assert obs.task_id == "task5_hard_alt"
+
+
+def test_reset_unknown_task():
+    env = DataCleaningEnv()
+    with pytest.raises(ValueError):
+        env.reset(task_id="nonexistent_task")
+
+
+def test_issues_detected_on_reset():
+    env = DataCleaningEnv()
+    obs = env.reset(task_id="task1_easy")
+    assert len(obs.issues_detected) > 0

tests/test_step.py

@@ -0,0 +1,63 @@
+from env.environment import DataCleaningEnv
+from env.models import Action
+
+
+def test_remove_duplicates_gives_positive_reward():
+    env = DataCleaningEnv()
+    env.reset(task_id="task1_easy")
+    result = env.step(Action(type="remove_duplicates"))
+    assert result.reward > 0
+    assert not result.done
+
+
+def test_fill_missing_median():
+    env = DataCleaningEnv()
+    env.reset(task_id="task1_easy")
+    result = env.step(Action(type="fill_missing", column="age", strategy="median"))
+    assert result.reward >= 0
+    assert env.current_df["age"].isna().sum() == 0
+
+
+def test_invalid_action_penalised():
+    env = DataCleaningEnv()
+    env.reset(task_id="task1_easy")
+    result = env.step(Action(type="fill_missing", column="nonexistent_col", strategy="mean"))
+    assert result.reward < 0
+    assert result.info["error"] is not None
+
+
+def test_submit_ends_episode():
+    env = DataCleaningEnv()
+    env.reset(task_id="task1_easy")
+    result = env.step(Action(type="submit"))
+    assert result.done
+    assert result.info["final_score"] >= 0.0
+
+
+def test_step_after_done_is_no_op():
+    env = DataCleaningEnv()
+    env.reset(task_id="task1_easy")
+    env.step(Action(type="submit"))
+    result = env.step(Action(type="remove_duplicates"))
+    assert result.done
+    assert 0.0 < result.reward < 1.0
+    assert result.reward == result.info["final_score"]
+
+
+def test_convert_type_datetime():
+    env = DataCleaningEnv()
+    env.reset(task_id="task2_medium")
+    result = env.step(Action(type="convert_type", column="date", target_type="datetime"))
+    assert result.reward > 0
+    sample = env.current_df["date"].dropna().iloc[0]
+    import re
+
+    assert re.match(r"\d{4}-\d{2}-\d{2}", str(sample))
+
+
+def test_clip_outliers():
+    env = DataCleaningEnv()
+    env.reset(task_id="task3_hard")
+    result = env.step(Action(type="clip_outliers", column="session_duration", lower=0.0, upper=1000.0))
+    assert result.reward > 0
+    assert env.current_df["session_duration"].max() <= 1000.0