Spaces:
Sleeping
Sleeping
Commit ·
0d7c387
0
Parent(s):
Clean project upload
Browse files- .dockerignore +9 -0
- .gitignore +12 -0
- CSV_DC_ENV +1 -0
- Dockerfile +52 -0
- README.md +114 -0
- __init__.py +21 -0
- client.py +38 -0
- inference.py +250 -0
- models.py +46 -0
- openenv.yaml +6 -0
- pyproject.toml +31 -0
- requirements.txt +6 -0
- server/__init__.py +1 -0
- server/app.py +33 -0
- server/csv_cleaning_environment.py +459 -0
- 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 |
+
}
|