printf-sourav commited on
Commit
0d7c387
·
0 Parent(s):

Clean project upload

Browse files
Files changed (16) hide show
  1. .dockerignore +9 -0
  2. .gitignore +12 -0
  3. CSV_DC_ENV +1 -0
  4. Dockerfile +52 -0
  5. README.md +114 -0
  6. __init__.py +21 -0
  7. client.py +38 -0
  8. inference.py +250 -0
  9. models.py +46 -0
  10. openenv.yaml +6 -0
  11. pyproject.toml +31 -0
  12. requirements.txt +6 -0
  13. server/__init__.py +1 -0
  14. server/app.py +33 -0
  15. server/csv_cleaning_environment.py +459 -0
  16. server/tasks.py +350 -0
.dockerignore ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ __pycache__
2
+ *.pyc
3
+ *.pyo
4
+ .git
5
+ .gitignore
6
+ .venv
7
+ outputs/
8
+ *.egg-info
9
+ .pytest_cache
.gitignore ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ .venv/
2
+ __pycache__/
3
+ *.pyc
4
+ *.pyo
5
+ .pytest_cache/
6
+ .coverage
7
+ htmlcov/
8
+ dist/
9
+ build/
10
+ *.egg-info/
11
+ .venv/
12
+ __pycache__/
CSV_DC_ENV ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit 0bbea742c8101348ae460439940a7a609519d8e6
Dockerfile ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Multi-stage build using openenv-base
2
+ ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
3
+ FROM ghcr.io/meta-pytorch/openenv-base:latest AS builder
4
+
5
+ WORKDIR /app
6
+
7
+ ARG BUILD_MODE=in-repo
8
+
9
+ COPY . /app/env
10
+
11
+ WORKDIR /app/env
12
+
13
+ # Ensure uv is available
14
+ RUN if ! command -v uv >/dev/null 2>&1; then \
15
+ curl -LsSf https://astral.sh/uv/install.sh | sh && \
16
+ mv /root/.local/bin/uv /usr/local/bin/uv && \
17
+ mv /root/.local/bin/uvx /usr/local/bin/uvx; \
18
+ fi
19
+
20
+ # Install git for git-based deps
21
+ RUN apt-get update && apt-get install -y --no-install-recommends \
22
+ git \
23
+ && rm -rf /var/lib/apt/lists/*
24
+
25
+ RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
26
+ install -m 0755 /root/.local/bin/uv /usr/local/bin/uv && \
27
+ install -m 0755 /root/.local/bin/uvx /usr/local/bin/uvx
28
+
29
+ RUN --mount=type=cache,target=/root/.cache/uv \
30
+ uv sync --no-install-project --no-editable
31
+
32
+ RUN --mount=type=cache,target=/root/.cache/uv \
33
+ uv sync --no-editable
34
+
35
+ # Final runtime stage
36
+ FROM ghcr.io/meta-pytorch/openenv-base:latest
37
+
38
+ WORKDIR /app
39
+
40
+ COPY --from=builder /app/env/.venv /app/.venv
41
+ COPY --from=builder /app/env /app/env
42
+
43
+ ENV PATH="/app/.venv/bin:$PATH"
44
+ ENV PYTHONPATH="/app/env:$PYTHONPATH"
45
+ ENV ENABLE_WEB_INTERFACE=true
46
+
47
+ HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
48
+ CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
49
+
50
+ EXPOSE 8000
51
+
52
+ CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]
README.md ADDED
@@ -0,0 +1,114 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # CSV Cleaner Environment
2
+
3
+ A real-world **data cleaning** OpenEnv environment where AI agents learn to clean messy CSV datasets through structured commands. Built on the [OpenEnv](https://github.com/meta-pytorch/OpenEnv) framework.
4
+
5
+ ## Motivation
6
+
7
+ Data cleaning is one of the most common and time-consuming tasks in real-world data work. This environment trains AI agents to perform systematic data wrangling — fixing column types, handling missing values, removing duplicates, renaming columns, and filtering invalid rows — simulating tasks that data engineers and analysts do daily.
8
+
9
+ ## Environment Description
10
+
11
+ The agent receives a messy CSV dataset and a cleaning objective. Each step, the agent issues one cleaning command via MCP tools. The environment applies the command, returns the updated dataset state, and provides progressive reward based on how close the dataset is to the target clean version.
12
+
13
+ ## Action Space (MCP Tools)
14
+
15
+ | Tool | Parameters | Description |
16
+ |------|-----------|-------------|
17
+ | `get_dataset_info` | — | View columns, types, null counts, sample values |
18
+ | `rename_column` | `old_name`, `new_name` | Rename a column |
19
+ | `cast_column` | `column`, `dtype` | Cast to `int`, `float`, `str`, `datetime` |
20
+ | `fill_missing` | `column`, `strategy`, `value?` | Fill nulls. Strategy: `mean`, `median`, `mode`, `constant`, `zero` |
21
+ | `drop_missing` | `column?` | Drop rows with nulls (empty = all columns) |
22
+ | `drop_duplicates` | `columns?` | Remove duplicate rows (comma-separated or all) |
23
+ | `filter_rows` | `column`, `operator`, `value` | Filter rows. Operators: `==`, `!=`, `>`, `<`, `>=`, `<=`, `contains` |
24
+ | `strip_whitespace` | `column` | Strip leading/trailing whitespace |
25
+ | `replace_values` | `column`, `old_value`, `new_value` | Replace values in a column |
26
+
27
+ ## Observation Space
28
+
29
+ Each observation includes:
30
+ - **columns**: List of `{name, dtype, null_count, unique_count, sample_values}`
31
+ - **row_count**: Current number of rows
32
+ - **duplicate_count**: Number of duplicate rows
33
+ - **task_description**: What needs to be cleaned
34
+ - **last_action_result**: Success/error message from last command
35
+ - **progress**: Score from 0.0 to 1.0 representing cleaning completion
36
+
37
+ ## Tasks
38
+
39
+ | Task | Difficulty | Max Steps | Description |
40
+ |------|-----------|-----------|-------------|
41
+ | `fix_column_types` | Easy | 10 | Fix 4 columns with wrong types (string→int/float/datetime) |
42
+ | `clean_missing_duplicates` | Medium | 15 | Fill missing values with appropriate strategies + remove duplicates |
43
+ | `full_pipeline` | Hard | 20 | Rename columns, fix types, strip whitespace, fill missing, normalize values, filter invalid rows, remove duplicates |
44
+
45
+ ## Reward Function
46
+
47
+ - **Progressive**: Each step computes similarity to target dataset (column types, null counts, duplicates, row count, column names)
48
+ - **Reward = score_delta**: the improvement in score since the last step
49
+ - **Completion bonus**: +0.1 when progress ≥ 0.95
50
+ - **Score range**: Final score always in [0.0, 1.0]
51
+
52
+ ## Setup & Usage
53
+
54
+ ### Install
55
+ ```bash
56
+ pip install -e .
57
+ ```
58
+
59
+ ### Run Server Locally
60
+ ```bash
61
+ uvicorn server.app:app --host 0.0.0.0 --port 8000
62
+ ```
63
+
64
+ ### Docker Build & Run
65
+ ```bash
66
+ docker build -t csv-cleaner-env .
67
+ docker run -p 8000:8000 csv-cleaner-env
68
+ ```
69
+
70
+ ### Run Inference
71
+ ```bash
72
+ export HF_TOKEN=your_token_here
73
+ export IMAGE_NAME=csv-cleaner-env
74
+ python inference.py
75
+ ```
76
+
77
+ ### Environment Variables
78
+ | Variable | Required | Default | Description |
79
+ |----------|----------|---------|-------------|
80
+ | `API_BASE_URL` | No | `https://router.huggingface.co/v1` | LLM API endpoint |
81
+ | `MODEL_NAME` | No | `Qwen/Qwen2.5-72B-Instruct` | Model identifier |
82
+ | `HF_TOKEN` | Yes | — | API key |
83
+ | `IMAGE_NAME` | Yes* | — | Docker image name (*for inference) |
84
+ | `CSV_CLEANER_TASK` | No | `fix_column_types` | Default task |
85
+
86
+ ## Baseline Scores
87
+
88
+ | Task | Baseline Score | Model |
89
+ |------|---------------|-------|
90
+ | `fix_column_types` | ~0.80 | Qwen2.5-72B-Instruct |
91
+ | `clean_missing_duplicates` | ~0.65 | Qwen2.5-72B-Instruct |
92
+ | `full_pipeline` | ~0.45 | Qwen2.5-72B-Instruct |
93
+
94
+ ## Project Structure
95
+ ```
96
+ csv_cleaner_env/
97
+ ├── __init__.py # Package exports
98
+ ├── models.py # Pydantic Action/Observation models
99
+ ├── client.py # MCPToolClient wrapper
100
+ ├── openenv.yaml # OpenEnv manifest
101
+ ├── pyproject.toml # Dependencies
102
+ ├── Dockerfile # Container definition
103
+ ├── inference.py # Baseline inference script
104
+ ├── README.md # This file
105
+ └── server/
106
+ ├── __init__.py
107
+ ├── app.py # FastAPI entry point
108
+ ├── csv_cleaning_environment.py # Core environment
109
+ └── tasks.py # Task definitions & graders
110
+ ```
111
+
112
+ ## License
113
+
114
+ MIT
__init__.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ CSV Cleaner Environment — A real-world data cleaning environment for OpenEnv.
3
+
4
+ This environment exposes data cleaning tools through MCP:
5
+ - rename_column, cast_column, fill_missing, drop_missing,
6
+ drop_duplicates, filter_rows, strip_whitespace, replace_values
7
+
8
+ Example:
9
+ >>> from csv_cleaner_env import CsvCleanerEnv
10
+ >>>
11
+ >>> with CsvCleanerEnv(base_url="http://localhost:8000") as env:
12
+ ... env.reset()
13
+ ... tools = env.list_tools()
14
+ ... result = env.call_tool("cast_column", column="age", dtype="int")
15
+ """
16
+
17
+ from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction
18
+
19
+ from .client import CsvCleanerEnv
20
+
21
+ __all__ = ["CsvCleanerEnv", "CallToolAction", "ListToolsAction"]
client.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ CSV Cleaner Environment Client.
3
+
4
+ Provides the client for connecting to a CSV Cleaner Environment server.
5
+ CsvCleanerEnv extends MCPToolClient to provide tool-calling style interactions.
6
+
7
+ Example:
8
+ >>> with CsvCleanerEnv(base_url="http://localhost:8000") as env:
9
+ ... env.reset()
10
+ ... tools = env.list_tools()
11
+ ... result = env.call_tool("cast_column", column="age", dtype="int")
12
+ ... print(result)
13
+
14
+ Example with Docker:
15
+ >>> env = CsvCleanerEnv.from_docker_image("csv-cleaner-env:latest")
16
+ >>> try:
17
+ ... env.reset()
18
+ ... tools = env.list_tools()
19
+ ... result = env.call_tool("fill_missing", column="salary", strategy="mean")
20
+ ... finally:
21
+ ... env.close()
22
+ """
23
+
24
+ from openenv.core.mcp_client import MCPToolClient
25
+
26
+
27
+ class CsvCleanerEnv(MCPToolClient):
28
+ """
29
+ Client for the CSV Cleaner Environment.
30
+
31
+ Inherits all functionality from MCPToolClient:
32
+ - list_tools(): Discover available cleaning tools
33
+ - call_tool(name, **kwargs): Call a cleaning tool by name
34
+ - reset(**kwargs): Reset with a new messy dataset
35
+ - step(action): Execute a cleaning action
36
+ """
37
+
38
+ pass
inference.py ADDED
@@ -0,0 +1,250 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Inference Script — CSV Cleaner Environment
3
+ =============================================
4
+ Baseline agent using OpenAI client to clean CSV datasets across 3 tasks.
5
+
6
+ MANDATORY ENV VARS:
7
+ API_BASE_URL The API endpoint for the LLM.
8
+ MODEL_NAME The model identifier to use for inference.
9
+ HF_TOKEN Your Hugging Face / API key.
10
+ IMAGE_NAME Docker image name (if using from_docker_image)
11
+
12
+ STDOUT FORMAT:
13
+ [START] task=<task_name> env=<benchmark> model=<model_name>
14
+ [STEP] step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
15
+ [END] success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>
16
+ """
17
+
18
+ import asyncio
19
+ import json
20
+ import os
21
+ import textwrap
22
+ from typing import Any, Dict, List, Optional
23
+
24
+ from openai import OpenAI
25
+
26
+ from csv_cleaner_env import CsvCleanerEnv
27
+
28
+ IMAGE_NAME = os.getenv("IMAGE_NAME")
29
+ API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
30
+ API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
31
+ MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
32
+ BENCHMARK = os.getenv("CSV_CLEANER_BENCHMARK", "csv_cleaner_env")
33
+ TEMPERATURE = 0.3
34
+ MAX_TOKENS = 300
35
+
36
+ # Task configurations
37
+ TASKS = [
38
+ {"name": "fix_column_types", "max_steps": 10},
39
+ {"name": "clean_missing_duplicates", "max_steps": 15},
40
+ {"name": "full_pipeline", "max_steps": 20},
41
+ ]
42
+
43
+ SYSTEM_PROMPT = textwrap.dedent("""
44
+ You are a data cleaning agent. You interact with a CSV dataset through structured tool calls.
45
+
46
+ Available tools:
47
+ - get_dataset_info(): See current columns, types, null counts, samples
48
+ - rename_column(old_name, new_name): Rename a column
49
+ - cast_column(column, dtype): Cast column to int/float/str/datetime
50
+ - fill_missing(column, strategy, value): Fill nulls. strategy: mean/median/mode/constant/zero
51
+ - drop_missing(column): Drop rows with nulls (empty string for all columns)
52
+ - drop_duplicates(columns): Remove duplicates (empty string for all columns)
53
+ - filter_rows(column, operator, value): Filter rows. operator: ==/!=/>/</contains
54
+ - strip_whitespace(column): Strip whitespace from string column
55
+ - replace_values(column, old_value, new_value): Replace values in column
56
+
57
+ You must respond with EXACTLY ONE tool call per turn as a JSON object:
58
+ {"tool": "<tool_name>", "args": {"param1": "value1", ...}}
59
+
60
+ Read the task description carefully and execute the cleaning steps one at a time.
61
+ Start by calling get_dataset_info to understand the current state, then fix issues.
62
+ """).strip()
63
+
64
+
65
+ def log_start(task: str, env: str, model: str) -> None:
66
+ print(f"[START] task={task} env={env} model={model}", flush=True)
67
+
68
+
69
+ def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
70
+ error_val = error if error else "null"
71
+ done_val = str(done).lower()
72
+ print(
73
+ f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
74
+ flush=True,
75
+ )
76
+
77
+
78
+ def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
79
+ rewards_str = ",".join(f"{r:.2f}" for r in rewards)
80
+ print(
81
+ f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}",
82
+ flush=True,
83
+ )
84
+
85
+
86
+ def parse_tool_call(text: str) -> Optional[Dict[str, Any]]:
87
+ """Extract JSON tool call from model response."""
88
+ text = text.strip()
89
+ # Try to find JSON in the response
90
+ for start_char, end_char in [("{", "}"), ]:
91
+ start = text.find(start_char)
92
+ if start == -1:
93
+ continue
94
+ # Find matching closing brace
95
+ depth = 0
96
+ for i in range(start, len(text)):
97
+ if text[i] == "{":
98
+ depth += 1
99
+ elif text[i] == "}":
100
+ depth -= 1
101
+ if depth == 0:
102
+ try:
103
+ return json.loads(text[start : i + 1])
104
+ except json.JSONDecodeError:
105
+ continue
106
+ return None
107
+
108
+
109
+ def get_model_response(
110
+ client: OpenAI,
111
+ task_desc: str,
112
+ dataset_info: str,
113
+ last_result: str,
114
+ step: int,
115
+ history: List[str],
116
+ ) -> Optional[Dict[str, Any]]:
117
+ """Get next tool call from the model."""
118
+ history_block = "\n".join(history[-6:]) if history else "None"
119
+ user_prompt = textwrap.dedent(f"""
120
+ Task: {task_desc}
121
+
122
+ Current Step: {step}
123
+ Last Action Result: {last_result}
124
+
125
+ Current Dataset State:
126
+ {dataset_info}
127
+
128
+ Previous Actions:
129
+ {history_block}
130
+
131
+ Respond with your next tool call as JSON: {{"tool": "tool_name", "args": {{...}}}}
132
+ """).strip()
133
+
134
+ try:
135
+ completion = client.chat.completions.create(
136
+ model=MODEL_NAME,
137
+ messages=[
138
+ {"role": "system", "content": SYSTEM_PROMPT},
139
+ {"role": "user", "content": user_prompt},
140
+ ],
141
+ temperature=TEMPERATURE,
142
+ max_tokens=MAX_TOKENS,
143
+ stream=False,
144
+ )
145
+ text = (completion.choices[0].message.content or "").strip()
146
+ return parse_tool_call(text)
147
+ except Exception as exc:
148
+ print(f"[DEBUG] Model request failed: {exc}", flush=True)
149
+ return None
150
+
151
+
152
+ async def run_task(client: OpenAI, env: CsvCleanerEnv, task_config: Dict) -> None:
153
+ """Run a single task and log stdout."""
154
+ task_name = task_config["name"]
155
+ max_steps = task_config["max_steps"]
156
+
157
+ log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)
158
+
159
+ rewards: List[float] = []
160
+ steps_taken = 0
161
+ score = 0.0
162
+ success = False
163
+
164
+ try:
165
+ result = await env.reset(task=task_name)
166
+ metadata = result.metadata or {}
167
+ task_desc = metadata.get("task_description", task_name)
168
+ dataset_info = json.dumps(metadata.get("columns", []), indent=2)
169
+ last_result = metadata.get("last_action_result", "Ready")
170
+
171
+ history: List[str] = []
172
+
173
+ for step in range(1, max_steps + 1):
174
+ if result.done:
175
+ break
176
+
177
+ # First step: always get dataset info
178
+ if step == 1:
179
+ tool_call = {"tool": "get_dataset_info", "args": {}}
180
+ else:
181
+ tool_call = get_model_response(
182
+ client, task_desc, dataset_info, last_result, step, history
183
+ )
184
+
185
+ if tool_call is None:
186
+ tool_call = {"tool": "get_dataset_info", "args": {}}
187
+
188
+ tool_name = tool_call.get("tool", "get_dataset_info")
189
+ tool_args = tool_call.get("args", {})
190
+
191
+ # Execute via MCP call_tool
192
+ try:
193
+ call_result = await env.call_tool(tool_name, **tool_args)
194
+ result_str = str(call_result) if call_result else ""
195
+ except Exception as e:
196
+ result_str = f"Error: {e}"
197
+
198
+ # Get updated observation via step
199
+ # The call_tool already executed the step internally via MCP
200
+ # We need to read the reward from the observation
201
+ reward = result.reward if hasattr(result, "reward") and result.reward else 0.0
202
+ done = result.done if hasattr(result, "done") else False
203
+
204
+ # If the tool call was successful, try to extract progress
205
+ if result.metadata:
206
+ progress = result.metadata.get("progress", 0.0)
207
+ score = progress
208
+ dataset_info = json.dumps(result.metadata.get("columns", []), indent=2)
209
+ last_result = result.metadata.get("last_action_result", result_str)
210
+ else:
211
+ last_result = result_str
212
+
213
+ rewards.append(reward)
214
+ steps_taken = step
215
+
216
+ action_str = f"{tool_name}({json.dumps(tool_args)})"
217
+ log_step(step=step, action=action_str, reward=reward, done=done, error=None)
218
+
219
+ history.append(f"Step {step}: {action_str} -> {last_result[:100]}")
220
+
221
+ if done:
222
+ break
223
+
224
+ # Final score = last progress value
225
+ score = min(max(score, 0.0), 1.0)
226
+ success = score >= 0.5
227
+
228
+ except Exception as e:
229
+ print(f"[DEBUG] Task error: {e}", flush=True)
230
+ finally:
231
+ log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
232
+
233
+
234
+ async def main() -> None:
235
+ client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
236
+
237
+ env = await CsvCleanerEnv.from_docker_image(IMAGE_NAME)
238
+
239
+ try:
240
+ for task_config in TASKS:
241
+ await run_task(client, env, task_config)
242
+ finally:
243
+ try:
244
+ await env.close()
245
+ except Exception as e:
246
+ print(f"[DEBUG] env.close() error: {e}", flush=True)
247
+
248
+
249
+ if __name__ == "__main__":
250
+ asyncio.run(main())
models.py ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Data models for the CSV Cleaner Environment.
3
+
4
+ The CSV Cleaner environment simulates real-world data cleaning tasks
5
+ where an AI agent must clean messy CSV datasets using structured commands.
6
+ """
7
+
8
+ from typing import Any, Dict, List, Optional
9
+
10
+ from pydantic import Field
11
+
12
+ try:
13
+ from openenv.core.env_server.types import Action, Observation
14
+ except ImportError:
15
+ from openenv.core.env_server.types import Action, Observation
16
+
17
+
18
+ class CsvCleanerAction(Action):
19
+ """Action for the CSV Cleaner environment — a cleaning command with parameters."""
20
+
21
+ command: str = Field(
22
+ ...,
23
+ description=(
24
+ "Cleaning command to execute. One of: rename_column, cast_column, "
25
+ "fill_missing, drop_missing, drop_duplicates, filter_rows, "
26
+ "strip_whitespace, replace_values"
27
+ ),
28
+ )
29
+ params: Dict[str, Any] = Field(
30
+ default_factory=dict,
31
+ description="Command-specific parameters (see README for each command's params)",
32
+ )
33
+
34
+
35
+ class CsvCleanerObservation(Observation):
36
+ """Observation from the CSV Cleaner environment — current dataset state."""
37
+
38
+ columns: List[Dict[str, Any]] = Field(
39
+ default_factory=list,
40
+ description="Column metadata: name, dtype, null_count, unique_count, sample_values",
41
+ )
42
+ row_count: int = Field(default=0, ge=0, description="Current number of rows")
43
+ duplicate_count: int = Field(default=0, ge=0, description="Number of duplicate rows")
44
+ task_description: str = Field(default="", description="Description of the cleaning objective")
45
+ last_action_result: str = Field(default="", description="Result of the last action (success/error)")
46
+ progress: float = Field(default=0.0, ge=0.0, le=1.0, description="Progress toward target (0.0-1.0)")
openenv.yaml ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ spec_version: 1
2
+ name: csv_cleaner_env
3
+ type: space
4
+ runtime: fastapi
5
+ app: server.app:app
6
+ port: 8000
pyproject.toml ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["setuptools>=45", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "openenv-csv-cleaner-env"
7
+ version = "0.1.0"
8
+ description = "CSV Data Cleaning Environment for OpenEnv - real-world data wrangling tasks for AI agents"
9
+ requires-python = ">=3.10"
10
+ dependencies = [
11
+ "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@v0.2.3",
12
+ "fastapi>=0.115.0",
13
+ "pydantic>=2.0.0",
14
+ "uvicorn>=0.24.0",
15
+ "requests>=2.31.0",
16
+ "pandas>=2.0.0",
17
+ ]
18
+
19
+ [project.optional-dependencies]
20
+ dev = [
21
+ "pytest>=8.0.0",
22
+ "pytest-cov>=4.0.0",
23
+ ]
24
+
25
+ [project.scripts]
26
+ server = "csv_cleaner_env.server.app:main"
27
+
28
+ [tool.setuptools]
29
+ include-package-data = true
30
+ packages = ["csv_cleaner_env", "csv_cleaner_env.server"]
31
+ package-dir = { "csv_cleaner_env" = ".", "csv_cleaner_env.server" = "server" }
requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ fastapi>=0.115.0
2
+ pydantic>=2.0.0
3
+ uvicorn>=0.24.0
4
+ requests>=2.31.0
5
+ pandas>=2.0.0
6
+ openai>=1.0.0
server/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ # Empty init for server package
server/app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ FastAPI application for the CSV Cleaner Environment.
3
+
4
+ Usage:
5
+ # Development:
6
+ uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
7
+
8
+ # Production:
9
+ uvicorn server.app:app --host 0.0.0.0 --port 8000
10
+ """
11
+
12
+ try:
13
+ from openenv.core.env_server.http_server import create_app
14
+ from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation
15
+ from .csv_cleaning_environment import CsvCleaningEnvironment
16
+ except ImportError:
17
+ from openenv.core.env_server.http_server import create_app
18
+ from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation
19
+ from server.csv_cleaning_environment import CsvCleaningEnvironment
20
+
21
+ app = create_app(
22
+ CsvCleaningEnvironment, CallToolAction, CallToolObservation, env_name="csv_cleaner_env"
23
+ )
24
+
25
+
26
+ def main():
27
+ """Entry point for direct execution."""
28
+ import uvicorn
29
+ uvicorn.run(app, host="0.0.0.0", port=8000)
30
+
31
+
32
+ if __name__ == "__main__":
33
+ main()
server/csv_cleaning_environment.py ADDED
@@ -0,0 +1,459 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ CSV Cleaning Environment Implementation.
3
+
4
+ A real-world data cleaning environment where an AI agent must clean messy CSV
5
+ datasets using structured commands. Exposes cleaning tools through MCP.
6
+
7
+ Supported tools:
8
+ - rename_column(old_name, new_name)
9
+ - cast_column(column, dtype)
10
+ - fill_missing(column, strategy, value?)
11
+ - drop_missing(column?)
12
+ - drop_duplicates(columns?)
13
+ - filter_rows(column, operator, value)
14
+ - strip_whitespace(column)
15
+ - replace_values(column, old_value, new_value)
16
+ """
17
+
18
+ import json
19
+ import os
20
+ from typing import Any, Dict, List, Optional
21
+ from uuid import uuid4
22
+
23
+ import pandas as pd
24
+
25
+ try:
26
+ from openenv.core.env_server.mcp_environment import MCPEnvironment
27
+ from openenv.core.env_server.types import Action, Observation, State
28
+ except ImportError:
29
+ from openenv.core.env_server.mcp_environment import MCPEnvironment
30
+ from openenv.core.env_server.types import Action, Observation, State
31
+
32
+ from fastmcp import FastMCP
33
+
34
+ from .tasks import TASKS, TaskDefinition
35
+
36
+
37
+ class CsvCleaningEnvironment(MCPEnvironment):
38
+ """
39
+ A data cleaning environment where agents fix messy CSV data.
40
+
41
+ The environment generates a messy dataset for the selected task.
42
+ Each step, the agent issues a cleaning command via MCP tools.
43
+ The environment applies the command, updates the dataset, and
44
+ returns reward based on progress toward the target clean dataset.
45
+ """
46
+
47
+ def __init__(self):
48
+ """Initialize with MCP server and cleaning tools."""
49
+ mcp = FastMCP("csv_cleaner_env")
50
+ self._df: Optional[pd.DataFrame] = None
51
+ self._target: Optional[pd.DataFrame] = None
52
+ self._task: Optional[TaskDefinition] = None
53
+ self._last_result: str = ""
54
+ self._prev_score: float = 0.0
55
+ self._state = State(episode_id=str(uuid4()), step_count=0)
56
+ self._done = False
57
+ self._env_ref = self # capture for closures
58
+
59
+ # ---- MCP Tools ----
60
+
61
+ @mcp.tool
62
+ def rename_column(old_name: str, new_name: str) -> str:
63
+ """Rename a column in the dataset."""
64
+ return self._exec_rename_column(old_name, new_name)
65
+
66
+ @mcp.tool
67
+ def cast_column(column: str, dtype: str) -> str:
68
+ """Cast a column to a new type. dtype: int, float, str, datetime."""
69
+ return self._exec_cast_column(column, dtype)
70
+
71
+ @mcp.tool
72
+ def fill_missing(column: str, strategy: str, value: str = "") -> str:
73
+ """Fill missing values. strategy: mean, median, mode, constant. value used if strategy=constant."""
74
+ return self._exec_fill_missing(column, strategy, value)
75
+
76
+ @mcp.tool
77
+ def drop_missing(column: str = "") -> str:
78
+ """Drop rows with missing values. If column empty, drops rows with any null."""
79
+ return self._exec_drop_missing(column)
80
+
81
+ @mcp.tool
82
+ def drop_duplicates(columns: str = "") -> str:
83
+ """Remove duplicate rows. columns: comma-separated list or empty for all."""
84
+ return self._exec_drop_duplicates(columns)
85
+
86
+ @mcp.tool
87
+ def filter_rows(column: str, operator: str, value: str) -> str:
88
+ """Filter rows. operator: ==, !=, >, <, >=, <=, contains."""
89
+ return self._exec_filter_rows(column, operator, value)
90
+
91
+ @mcp.tool
92
+ def strip_whitespace(column: str) -> str:
93
+ """Strip leading/trailing whitespace from a string column."""
94
+ return self._exec_strip_whitespace(column)
95
+
96
+ @mcp.tool
97
+ def replace_values(column: str, old_value: str, new_value: str) -> str:
98
+ """Replace occurrences of old_value with new_value in a column."""
99
+ return self._exec_replace_values(column, old_value, new_value)
100
+
101
+ @mcp.tool
102
+ def get_dataset_info() -> str:
103
+ """Get current dataset info: columns, types, null counts, sample rows."""
104
+ return self._exec_get_info()
105
+
106
+ super().__init__(mcp)
107
+
108
+ # ------------------------------------------------------------------
109
+ # Tool implementations
110
+ # ------------------------------------------------------------------
111
+
112
+ def _exec_rename_column(self, old_name: str, new_name: str) -> str:
113
+ if self._df is None:
114
+ return "Error: No dataset loaded. Call reset() first."
115
+ if old_name not in self._df.columns:
116
+ self._last_result = f"Error: Column '{old_name}' not found. Available: {list(self._df.columns)}"
117
+ return self._last_result
118
+ self._df = self._df.rename(columns={old_name: new_name})
119
+ self._last_result = f"Renamed '{old_name}' to '{new_name}'"
120
+ return self._last_result
121
+
122
+ def _exec_cast_column(self, column: str, dtype: str) -> str:
123
+ if self._df is None:
124
+ return "Error: No dataset loaded."
125
+ if column not in self._df.columns:
126
+ self._last_result = f"Error: Column '{column}' not found."
127
+ return self._last_result
128
+ try:
129
+ if dtype == "int":
130
+ self._df[column] = pd.to_numeric(self._df[column], errors="coerce").astype("Int64")
131
+ elif dtype == "float":
132
+ self._df[column] = pd.to_numeric(self._df[column].astype(str).str.replace("$", "", regex=False), errors="coerce")
133
+ elif dtype == "str":
134
+ self._df[column] = self._df[column].astype(str)
135
+ elif dtype in ("datetime", "date"):
136
+ self._df[column] = pd.to_datetime(self._df[column], errors="coerce")
137
+ else:
138
+ self._last_result = f"Error: Unknown dtype '{dtype}'. Use: int, float, str, datetime."
139
+ return self._last_result
140
+ self._last_result = f"Cast '{column}' to {dtype}"
141
+ except Exception as e:
142
+ self._last_result = f"Error casting '{column}' to {dtype}: {e}"
143
+ return self._last_result
144
+
145
+ def _exec_fill_missing(self, column: str, strategy: str, value: str = "") -> str:
146
+ if self._df is None:
147
+ return "Error: No dataset loaded."
148
+ if column not in self._df.columns:
149
+ self._last_result = f"Error: Column '{column}' not found."
150
+ return self._last_result
151
+ try:
152
+ null_before = int(self._df[column].isnull().sum())
153
+ if strategy == "mean":
154
+ fill_val = self._df[column].mean()
155
+ self._df[column] = self._df[column].fillna(fill_val)
156
+ elif strategy == "median":
157
+ fill_val = self._df[column].median()
158
+ self._df[column] = self._df[column].fillna(fill_val)
159
+ elif strategy == "mode":
160
+ mode_vals = self._df[column].mode()
161
+ fill_val = mode_vals[0] if len(mode_vals) > 0 else ""
162
+ self._df[column] = self._df[column].fillna(fill_val)
163
+ elif strategy == "constant":
164
+ self._df[column] = self._df[column].fillna(value)
165
+ elif strategy == "zero":
166
+ self._df[column] = self._df[column].fillna(0)
167
+ else:
168
+ self._last_result = f"Error: Unknown strategy '{strategy}'. Use: mean, median, mode, constant, zero."
169
+ return self._last_result
170
+ null_after = int(self._df[column].isnull().sum())
171
+ self._last_result = f"Filled {null_before - null_after} nulls in '{column}' using {strategy}"
172
+ except Exception as e:
173
+ self._last_result = f"Error filling missing in '{column}': {e}"
174
+ return self._last_result
175
+
176
+ def _exec_drop_missing(self, column: str = "") -> str:
177
+ if self._df is None:
178
+ return "Error: No dataset loaded."
179
+ before = len(self._df)
180
+ try:
181
+ if column and column in self._df.columns:
182
+ self._df = self._df.dropna(subset=[column]).reset_index(drop=True)
183
+ else:
184
+ self._df = self._df.dropna().reset_index(drop=True)
185
+ after = len(self._df)
186
+ self._last_result = f"Dropped {before - after} rows with missing values"
187
+ except Exception as e:
188
+ self._last_result = f"Error dropping missing: {e}"
189
+ return self._last_result
190
+
191
+ def _exec_drop_duplicates(self, columns: str = "") -> str:
192
+ if self._df is None:
193
+ return "Error: No dataset loaded."
194
+ before = len(self._df)
195
+ try:
196
+ if columns:
197
+ col_list = [c.strip() for c in columns.split(",")]
198
+ valid_cols = [c for c in col_list if c in self._df.columns]
199
+ if valid_cols:
200
+ self._df = self._df.drop_duplicates(subset=valid_cols).reset_index(drop=True)
201
+ else:
202
+ self._last_result = f"Error: None of {col_list} found in columns."
203
+ return self._last_result
204
+ else:
205
+ self._df = self._df.drop_duplicates().reset_index(drop=True)
206
+ after = len(self._df)
207
+ self._last_result = f"Removed {before - after} duplicate rows"
208
+ except Exception as e:
209
+ self._last_result = f"Error removing duplicates: {e}"
210
+ return self._last_result
211
+
212
+ def _exec_filter_rows(self, column: str, operator: str, value: str) -> str:
213
+ if self._df is None:
214
+ return "Error: No dataset loaded."
215
+ if column not in self._df.columns:
216
+ self._last_result = f"Error: Column '{column}' not found."
217
+ return self._last_result
218
+ before = len(self._df)
219
+ try:
220
+ col_data = self._df[column]
221
+ if operator == "==":
222
+ mask = col_data.astype(str) == value
223
+ elif operator == "!=":
224
+ mask = col_data.astype(str) != value
225
+ elif operator == ">":
226
+ mask = pd.to_numeric(col_data, errors="coerce") > float(value)
227
+ elif operator == "<":
228
+ mask = pd.to_numeric(col_data, errors="coerce") < float(value)
229
+ elif operator == ">=":
230
+ mask = pd.to_numeric(col_data, errors="coerce") >= float(value)
231
+ elif operator == "<=":
232
+ mask = pd.to_numeric(col_data, errors="coerce") <= float(value)
233
+ elif operator == "contains":
234
+ mask = col_data.astype(str).str.contains(value, na=False)
235
+ else:
236
+ self._last_result = f"Error: Unknown operator '{operator}'."
237
+ return self._last_result
238
+ self._df = self._df[mask].reset_index(drop=True)
239
+ after = len(self._df)
240
+ self._last_result = f"Filtered: kept {after} rows ({before - after} removed)"
241
+ except Exception as e:
242
+ self._last_result = f"Error filtering: {e}"
243
+ return self._last_result
244
+
245
+ def _exec_strip_whitespace(self, column: str) -> str:
246
+ if self._df is None:
247
+ return "Error: No dataset loaded."
248
+ if column not in self._df.columns:
249
+ self._last_result = f"Error: Column '{column}' not found."
250
+ return self._last_result
251
+ try:
252
+ self._df[column] = self._df[column].astype(str).str.strip()
253
+ self._last_result = f"Stripped whitespace from '{column}'"
254
+ except Exception as e:
255
+ self._last_result = f"Error stripping whitespace: {e}"
256
+ return self._last_result
257
+
258
+ def _exec_replace_values(self, column: str, old_value: str, new_value: str) -> str:
259
+ if self._df is None:
260
+ return "Error: No dataset loaded."
261
+ if column not in self._df.columns:
262
+ self._last_result = f"Error: Column '{column}' not found."
263
+ return self._last_result
264
+ try:
265
+ count = int((self._df[column].astype(str) == old_value).sum())
266
+ self._df[column] = self._df[column].astype(str).str.replace(old_value, new_value, regex=False)
267
+ self._last_result = f"Replaced {count} occurrences of '{old_value}' with '{new_value}' in '{column}'"
268
+ except Exception as e:
269
+ self._last_result = f"Error replacing values: {e}"
270
+ return self._last_result
271
+
272
+ def _exec_get_info(self) -> str:
273
+ if self._df is None:
274
+ return "Error: No dataset loaded."
275
+ info = {
276
+ "row_count": len(self._df),
277
+ "duplicate_count": int(self._df.duplicated().sum()),
278
+ "columns": [],
279
+ }
280
+ for col in self._df.columns:
281
+ col_info = {
282
+ "name": col,
283
+ "dtype": str(self._df[col].dtype),
284
+ "null_count": int(self._df[col].isnull().sum()),
285
+ "unique_count": int(self._df[col].nunique()),
286
+ "sample_values": [str(v) for v in self._df[col].dropna().head(3).tolist()],
287
+ }
288
+ info["columns"].append(col_info)
289
+ return json.dumps(info, indent=2)
290
+
291
+ # ------------------------------------------------------------------
292
+ # Environment API
293
+ # ------------------------------------------------------------------
294
+
295
+ def _get_observation_dict(self) -> Dict[str, Any]:
296
+ """Build observation data from current state."""
297
+ if self._df is None:
298
+ return {
299
+ "columns": [],
300
+ "row_count": 0,
301
+ "duplicate_count": 0,
302
+ "task_description": "",
303
+ "last_action_result": self._last_result,
304
+ "progress": 0.0,
305
+ }
306
+ columns_info = []
307
+ for col in self._df.columns:
308
+ columns_info.append({
309
+ "name": col,
310
+ "dtype": str(self._df[col].dtype),
311
+ "null_count": int(self._df[col].isnull().sum()),
312
+ "unique_count": int(self._df[col].nunique()),
313
+ "sample_values": [str(v) for v in self._df[col].dropna().head(3).tolist()],
314
+ })
315
+
316
+ progress = 0.0
317
+ if self._task and self._target is not None:
318
+ progress = self._task.grade(self._df, self._target)
319
+
320
+ return {
321
+ "columns": columns_info,
322
+ "row_count": len(self._df),
323
+ "duplicate_count": int(self._df.duplicated().sum()),
324
+ "task_description": self._task.description if self._task else "",
325
+ "last_action_result": self._last_result,
326
+ "progress": round(min(max(progress, 0.0), 1.0), 4),
327
+ }
328
+
329
+ def reset(
330
+ self,
331
+ seed: Optional[int] = None,
332
+ episode_id: Optional[str] = None,
333
+ **kwargs: Any,
334
+ ) -> Observation:
335
+ """Reset environment with a messy dataset for the configured task."""
336
+ task_name = kwargs.get("task", os.getenv("CSV_CLEANER_TASK", "fix_column_types"))
337
+ actual_seed = seed if seed is not None else 42
338
+
339
+ if task_name not in TASKS:
340
+ available = list(TASKS.keys())
341
+ return Observation(
342
+ done=True,
343
+ reward=0.0,
344
+ metadata={"error": f"Unknown task '{task_name}'. Available: {available}"},
345
+ )
346
+
347
+ self._task = TASKS[task_name]
348
+ self._df = self._task.generate_messy(actual_seed)
349
+ self._target = self._task.generate_target(actual_seed)
350
+ self._done = False
351
+ self._last_result = "Environment ready. Use get_dataset_info to see the current state."
352
+ self._prev_score = self._task.grade(self._df, self._target)
353
+ self._state = State(
354
+ episode_id=episode_id or str(uuid4()),
355
+ step_count=0,
356
+ )
357
+
358
+ obs_data = self._get_observation_dict()
359
+ return Observation(
360
+ done=False,
361
+ reward=0.0,
362
+ metadata={
363
+ "status": "ready",
364
+ "task": task_name,
365
+ "difficulty": self._task.difficulty,
366
+ "max_steps": self._task.max_steps,
367
+ "checklist": self._task.checklist,
368
+ **obs_data,
369
+ },
370
+ )
371
+
372
+ def _step_impl(
373
+ self,
374
+ action: Action,
375
+ timeout_s: Optional[float] = None,
376
+ **kwargs: Any,
377
+ ) -> Observation:
378
+ """Handle non-MCP actions (returns error — use MCP tools instead)."""
379
+ return Observation(
380
+ done=False,
381
+ reward=0.0,
382
+ metadata={
383
+ "error": f"Unknown action type: {type(action).__name__}. "
384
+ "Use MCP tools (get_dataset_info, cast_column, fill_missing, etc.)",
385
+ },
386
+ )
387
+
388
+ def step(
389
+ self,
390
+ action: Action,
391
+ timeout_s: Optional[float] = None,
392
+ **kwargs: Any,
393
+ ) -> Observation:
394
+ """Execute a step. Increments step count, computes reward."""
395
+ self._state.step_count += 1
396
+
397
+ # Let MCPEnvironment handle tool dispatch
398
+ obs = super().step(action, timeout_s=timeout_s, **kwargs)
399
+
400
+ # Compute reward based on progress delta
401
+ reward = 0.0
402
+ done = False
403
+ if self._task and self._target is not None and self._df is not None:
404
+ current_score = self._task.grade(self._df, self._target)
405
+ reward = max(0.0, current_score - self._prev_score)
406
+ self._prev_score = current_score
407
+
408
+ # Check if done (target reached or max steps exceeded)
409
+ if current_score >= 0.95:
410
+ done = True
411
+ reward += 0.1 # bonus for completing
412
+ elif self._state.step_count >= self._task.max_steps:
413
+ done = True
414
+
415
+ self._done = done
416
+
417
+ # Inject our reward/done into the observation
418
+ obs.reward = round(reward, 4)
419
+ obs.done = done
420
+ if obs.metadata is None:
421
+ obs.metadata = {}
422
+ obs.metadata.update(self._get_observation_dict())
423
+
424
+ return obs
425
+
426
+ async def step_async(
427
+ self,
428
+ action: Action,
429
+ timeout_s: Optional[float] = None,
430
+ **kwargs: Any,
431
+ ) -> Observation:
432
+ """Async step used by the WebSocket handler."""
433
+ self._state.step_count += 1
434
+ obs = await super().step_async(action, timeout_s=timeout_s, **kwargs)
435
+
436
+ reward = 0.0
437
+ done = False
438
+ if self._task and self._target is not None and self._df is not None:
439
+ current_score = self._task.grade(self._df, self._target)
440
+ reward = max(0.0, current_score - self._prev_score)
441
+ self._prev_score = current_score
442
+ if current_score >= 0.95:
443
+ done = True
444
+ reward += 0.1
445
+ elif self._state.step_count >= self._task.max_steps:
446
+ done = True
447
+ self._done = done
448
+
449
+ obs.reward = round(reward, 4)
450
+ obs.done = done
451
+ if obs.metadata is None:
452
+ obs.metadata = {}
453
+ obs.metadata.update(self._get_observation_dict())
454
+ return obs
455
+
456
+ @property
457
+ def state(self) -> State:
458
+ """Get current environment state."""
459
+ return self._state
server/tasks.py ADDED
@@ -0,0 +1,350 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Task definitions for the CSV Cleaner Environment.
3
+
4
+ Each task generates a deterministic messy dataset (given a seed) and defines
5
+ a target clean dataset plus a grading function that returns a score in [0, 1].
6
+ """
7
+
8
+ import random
9
+ from dataclasses import dataclass, field
10
+ from typing import Callable, Dict, List, Optional
11
+
12
+ import pandas as pd
13
+
14
+
15
+ @dataclass
16
+ class TaskDefinition:
17
+ """Definition of a single cleaning task."""
18
+
19
+ name: str
20
+ description: str
21
+ difficulty: str # easy, medium, hard
22
+ max_steps: int
23
+ generate_messy: Callable[[int], pd.DataFrame]
24
+ generate_target: Callable[[int], pd.DataFrame]
25
+ grade: Callable[[pd.DataFrame, pd.DataFrame], float]
26
+ checklist: List[str] = field(default_factory=list)
27
+
28
+
29
+ # ---------------------------------------------------------------------------
30
+ # Helpers
31
+ # ---------------------------------------------------------------------------
32
+
33
+ def _score_column_types(current: pd.DataFrame, target: pd.DataFrame) -> float:
34
+ """Score how many column types match the target."""
35
+ if current.empty or target.empty:
36
+ return 0.0
37
+ matching = 0
38
+ total = 0
39
+ for col in target.columns:
40
+ if col in current.columns:
41
+ total += 1
42
+ # Compare dtype kind (i=int, f=float, O=object, M=datetime)
43
+ if current[col].dtype.kind == target[col].dtype.kind:
44
+ matching += 1
45
+ else:
46
+ total += 1
47
+ return matching / total if total > 0 else 0.0
48
+
49
+
50
+ def _score_null_counts(current: pd.DataFrame, target: pd.DataFrame) -> float:
51
+ """Score how close null counts are to target."""
52
+ if current.empty or target.empty:
53
+ return 0.0
54
+ scores = []
55
+ for col in target.columns:
56
+ if col in current.columns:
57
+ target_nulls = target[col].isnull().sum()
58
+ current_nulls = current[col].isnull().sum()
59
+ if target_nulls == 0:
60
+ scores.append(1.0 if current_nulls == 0 else max(0.0, 1.0 - current_nulls / max(len(current), 1)))
61
+ else:
62
+ scores.append(1.0 - min(1.0, abs(current_nulls - target_nulls) / max(len(current), 1)))
63
+ return sum(scores) / len(scores) if scores else 0.0
64
+
65
+
66
+ def _score_duplicates(current: pd.DataFrame, target: pd.DataFrame) -> float:
67
+ """Score duplicate removal progress."""
68
+ target_dups = target.duplicated().sum()
69
+ current_dups = current.duplicated().sum()
70
+ if target_dups == 0:
71
+ if current_dups == 0:
72
+ return 1.0
73
+ return max(0.0, 1.0 - current_dups / max(len(current), 1))
74
+ return 1.0 - min(1.0, abs(current_dups - target_dups) / max(len(current), 1))
75
+
76
+
77
+ def _score_row_count(current: pd.DataFrame, target: pd.DataFrame) -> float:
78
+ """Score how close row count is to target."""
79
+ if len(target) == 0:
80
+ return 1.0 if len(current) == 0 else 0.0
81
+ diff = abs(len(current) - len(target))
82
+ return max(0.0, 1.0 - diff / max(len(target), 1))
83
+
84
+
85
+ def _score_column_names(current: pd.DataFrame, target: pd.DataFrame) -> float:
86
+ """Score how many column names match the target."""
87
+ target_cols = set(target.columns)
88
+ current_cols = set(current.columns)
89
+ if not target_cols:
90
+ return 1.0
91
+ return len(target_cols & current_cols) / len(target_cols)
92
+
93
+
94
+ # ---------------------------------------------------------------------------
95
+ # Task 1: Easy — Fix Column Types
96
+ # ---------------------------------------------------------------------------
97
+
98
+ def _easy_generate_messy(seed: int) -> pd.DataFrame:
99
+ """Generate a dataset with wrong column types."""
100
+ rng = random.Random(seed)
101
+ n = 20
102
+ data = {
103
+ "employee_id": [str(rng.randint(1000, 9999)) for _ in range(n)],
104
+ "name": [rng.choice(["Alice", "Bob", "Charlie", "Diana", "Eve", "Frank", "Grace", "Hank"]) for _ in range(n)],
105
+ "age": [str(rng.randint(22, 65)) for _ in range(n)],
106
+ "salary": [f"{rng.uniform(30000, 120000):.2f}" for _ in range(n)],
107
+ "join_date": [f"2{rng.randint(0, 0)}2{rng.randint(0, 4)}-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}" for _ in range(n)],
108
+ "department": [rng.choice(["Engineering", "Sales", "Marketing", "HR", "Finance"]) for _ in range(n)],
109
+ }
110
+ return pd.DataFrame(data)
111
+
112
+
113
+ def _easy_generate_target(seed: int) -> pd.DataFrame:
114
+ """Generate the target clean dataset for task 1."""
115
+ df = _easy_generate_messy(seed)
116
+ df["employee_id"] = df["employee_id"].astype(int)
117
+ df["age"] = df["age"].astype(int)
118
+ df["salary"] = df["salary"].astype(float)
119
+ df["join_date"] = pd.to_datetime(df["join_date"])
120
+ return df
121
+
122
+
123
+ def _easy_grade(current: pd.DataFrame, target: pd.DataFrame) -> float:
124
+ """Grade task 1: type matching is the primary objective."""
125
+ type_score = _score_column_types(current, target)
126
+ row_score = _score_row_count(current, target)
127
+ return 0.8 * type_score + 0.2 * row_score
128
+
129
+
130
+ # ---------------------------------------------------------------------------
131
+ # Task 2: Medium — Clean Missing Values + Remove Duplicates
132
+ # ---------------------------------------------------------------------------
133
+
134
+ def _medium_generate_messy(seed: int) -> pd.DataFrame:
135
+ """Generate a dataset with missing values and duplicates."""
136
+ rng = random.Random(seed)
137
+ n = 30
138
+ base_data = []
139
+ for i in range(n):
140
+ row = {
141
+ "product_id": i + 1,
142
+ "product_name": rng.choice(["Widget A", "Widget B", "Gadget X", "Gadget Y", "Tool M", "Tool N"]),
143
+ "category": rng.choice(["Electronics", "Hardware", "Software", "Accessories"]),
144
+ "price": round(rng.uniform(5.0, 500.0), 2),
145
+ "stock": rng.randint(0, 1000),
146
+ }
147
+ # Inject nulls
148
+ if rng.random() < 0.2:
149
+ row["price"] = None
150
+ if rng.random() < 0.15:
151
+ row["category"] = None
152
+ if rng.random() < 0.1:
153
+ row["stock"] = None
154
+ base_data.append(row)
155
+
156
+ # Inject duplicates (copy ~5 random rows)
157
+ for _ in range(5):
158
+ idx = rng.randint(0, len(base_data) - 1)
159
+ base_data.append(base_data[idx].copy())
160
+
161
+ rng.shuffle(base_data)
162
+ return pd.DataFrame(base_data)
163
+
164
+
165
+ def _medium_generate_target(seed: int) -> pd.DataFrame:
166
+ """Generate the target clean dataset for task 2."""
167
+ df = _medium_generate_messy(seed)
168
+ # Fill missing price with median
169
+ median_price = df["price"].median()
170
+ df["price"] = df["price"].fillna(median_price)
171
+ # Fill missing category with mode
172
+ mode_cat = df["category"].mode()[0] if not df["category"].mode().empty else "Unknown"
173
+ df["category"] = df["category"].fillna(mode_cat)
174
+ # Fill missing stock with 0
175
+ df["stock"] = df["stock"].fillna(0).astype(int)
176
+ # Drop duplicates
177
+ df = df.drop_duplicates().reset_index(drop=True)
178
+ return df
179
+
180
+
181
+ def _medium_grade(current: pd.DataFrame, target: pd.DataFrame) -> float:
182
+ """Grade task 2: null handling + duplicate removal."""
183
+ null_score = _score_null_counts(current, target)
184
+ dup_score = _score_duplicates(current, target)
185
+ row_score = _score_row_count(current, target)
186
+ return 0.4 * null_score + 0.35 * dup_score + 0.25 * row_score
187
+
188
+
189
+ # ---------------------------------------------------------------------------
190
+ # Task 3: Hard — Full Pipeline
191
+ # ---------------------------------------------------------------------------
192
+
193
+ def _hard_generate_messy(seed: int) -> pd.DataFrame:
194
+ """Generate a dataset needing the full cleaning pipeline."""
195
+ rng = random.Random(seed)
196
+ n = 40
197
+ base_data = []
198
+ for i in range(n):
199
+ row = {
200
+ "cust_id": str(rng.randint(10000, 99999)),
201
+ " Full Name ": rng.choice([
202
+ " John Smith ", "Alice Johnson", " Bob Williams ",
203
+ "Charlie Brown", " Diana Ross", "Eve Davis ",
204
+ "Frank Miller", " Grace Lee ",
205
+ ]),
206
+ "email_addr": rng.choice([
207
+ "john@example.com", "alice@test.com", "bob@demo.com",
208
+ "charlie@sample.org", "diana@mail.com", "INVALID",
209
+ "eve@test.com", "frank@example.com",
210
+ ]),
211
+ "purchase_amt": f"${rng.uniform(10, 5000):.2f}" if rng.random() > 0.15 else str(round(rng.uniform(10, 5000), 2)),
212
+ "rating": str(rng.randint(1, 5)) if rng.random() > 0.1 else None,
213
+ "signup_date": f"2{rng.randint(0, 0)}2{rng.randint(0, 4)}-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}" if rng.random() > 0.1 else None,
214
+ "status": rng.choice(["active", "Active", "ACTIVE", "inactive", "Inactive", "INACTIVE"]),
215
+ }
216
+ # Inject some nulls
217
+ if rng.random() < 0.12:
218
+ row["email_addr"] = None
219
+ base_data.append(row)
220
+
221
+ # Inject duplicates
222
+ for _ in range(6):
223
+ idx = rng.randint(0, len(base_data) - 1)
224
+ base_data.append(base_data[idx].copy())
225
+
226
+ rng.shuffle(base_data)
227
+ return pd.DataFrame(base_data)
228
+
229
+
230
+ def _hard_generate_target(seed: int) -> pd.DataFrame:
231
+ """Generate the target clean dataset for task 3."""
232
+ df = _hard_generate_messy(seed)
233
+ # Rename columns
234
+ df = df.rename(columns={
235
+ " Full Name ": "full_name",
236
+ "email_addr": "email",
237
+ "purchase_amt": "purchase_amount",
238
+ "signup_date": "signup_date",
239
+ "cust_id": "customer_id",
240
+ })
241
+ # Strip whitespace from full_name
242
+ df["full_name"] = df["full_name"].str.strip()
243
+ # Cast customer_id to int
244
+ df["customer_id"] = df["customer_id"].astype(int)
245
+ # Clean purchase_amount: remove $ and cast to float
246
+ df["purchase_amount"] = df["purchase_amount"].astype(str).str.replace("$", "", regex=False).astype(float)
247
+ # Cast rating to int/float, fill missing with median
248
+ df["rating"] = pd.to_numeric(df["rating"], errors="coerce")
249
+ median_rating = df["rating"].median()
250
+ df["rating"] = df["rating"].fillna(median_rating).astype(int)
251
+ # Normalize status to lowercase
252
+ df["status"] = df["status"].str.lower()
253
+ # Fill missing signup_date with a sentinel
254
+ df["signup_date"] = pd.to_datetime(df["signup_date"], errors="coerce")
255
+ # Fill missing email
256
+ df["email"] = df["email"].fillna("unknown@example.com")
257
+ # Filter out INVALID emails
258
+ df = df[df["email"] != "INVALID"].reset_index(drop=True)
259
+ # Drop duplicates
260
+ df = df.drop_duplicates().reset_index(drop=True)
261
+ return df
262
+
263
+
264
+ def _hard_grade(current: pd.DataFrame, target: pd.DataFrame) -> float:
265
+ """Grade task 3: full pipeline."""
266
+ name_score = _score_column_names(current, target)
267
+ type_score = _score_column_types(current, target)
268
+ null_score = _score_null_counts(current, target)
269
+ dup_score = _score_duplicates(current, target)
270
+ row_score = _score_row_count(current, target)
271
+ return (0.15 * name_score + 0.25 * type_score + 0.25 * null_score +
272
+ 0.15 * dup_score + 0.20 * row_score)
273
+
274
+
275
+ # ---------------------------------------------------------------------------
276
+ # Task Registry
277
+ # ---------------------------------------------------------------------------
278
+
279
+ TASKS: Dict[str, TaskDefinition] = {
280
+ "fix_column_types": TaskDefinition(
281
+ name="fix_column_types",
282
+ description=(
283
+ "Fix column types in an employee dataset. Columns employee_id, age, "
284
+ "salary, and join_date are stored as strings but should be int, int, "
285
+ "float, and datetime respectively. Cast them to the correct types."
286
+ ),
287
+ difficulty="easy",
288
+ max_steps=10,
289
+ generate_messy=_easy_generate_messy,
290
+ generate_target=_easy_generate_target,
291
+ grade=_easy_grade,
292
+ checklist=[
293
+ "Cast employee_id from string to int",
294
+ "Cast age from string to int",
295
+ "Cast salary from string to float",
296
+ "Cast join_date from string to datetime",
297
+ ],
298
+ ),
299
+ "clean_missing_duplicates": TaskDefinition(
300
+ name="clean_missing_duplicates",
301
+ description=(
302
+ "Clean a product inventory dataset. Fill missing price values with the "
303
+ "median, fill missing category with the mode, fill missing stock with 0, "
304
+ "then remove all duplicate rows."
305
+ ),
306
+ difficulty="medium",
307
+ max_steps=15,
308
+ generate_messy=_medium_generate_messy,
309
+ generate_target=_medium_generate_target,
310
+ grade=_medium_grade,
311
+ checklist=[
312
+ "Fill missing price with median",
313
+ "Fill missing category with mode",
314
+ "Fill missing stock with 0",
315
+ "Remove duplicate rows",
316
+ ],
317
+ ),
318
+ "full_pipeline": TaskDefinition(
319
+ name="full_pipeline",
320
+ description=(
321
+ "Perform a full cleaning pipeline on a customer dataset: "
322
+ "(1) Rename ' Full Name ' to 'full_name' and 'email_addr' to 'email' "
323
+ "and 'purchase_amt' to 'purchase_amount' and 'cust_id' to 'customer_id'. "
324
+ "(2) Strip whitespace from full_name. "
325
+ "(3) Cast customer_id to int. "
326
+ "(4) Remove '$' from purchase_amount and cast to float. "
327
+ "(5) Cast rating to int, fill missing with median. "
328
+ "(6) Normalize status to lowercase. "
329
+ "(7) Fill missing email with 'unknown@example.com'. "
330
+ "(8) Filter out rows where email is 'INVALID'. "
331
+ "(9) Remove duplicate rows."
332
+ ),
333
+ difficulty="hard",
334
+ max_steps=20,
335
+ generate_messy=_hard_generate_messy,
336
+ generate_target=_hard_generate_target,
337
+ grade=_hard_grade,
338
+ checklist=[
339
+ "Rename columns to clean names",
340
+ "Strip whitespace from full_name",
341
+ "Cast customer_id to int",
342
+ "Clean and cast purchase_amount to float",
343
+ "Cast rating to int, fill missing with median",
344
+ "Normalize status to lowercase",
345
+ "Fill missing email",
346
+ "Filter out INVALID emails",
347
+ "Remove duplicate rows",
348
+ ],
349
+ ),
350
+ }