Initial commit: OpenEnv Data Wrangler

.gitignore

@@ -0,0 +1,29 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+.env
+.venv
+env/
+venv/
+ENV/
+
+# OpenEnv generated
+uv.lock

Dockerfile

@@ -0,0 +1,80 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Multi-stage build using openenv-base
+# This Dockerfile is flexible and works for both:
+# - In-repo environments (with local OpenEnv sources)
+# - Standalone environments (with openenv from PyPI/Git)
+# The build script (openenv build) handles context detection and sets appropriate build args.
+
+ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+# Ensure git is available (required for installing dependencies from VCS)
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git && \
+    rm -rf /var/lib/apt/lists/*
+
+# Build argument to control whether we're building standalone or in-repo
+ARG BUILD_MODE=in-repo
+ARG ENV_NAME=data_wrangler
+
+# Copy environment code (always at root of build context)
+COPY . /app/env
+
+# For in-repo builds, openenv is already vendored in the build context
+# For standalone builds, openenv will be installed via pyproject.toml
+WORKDIR /app/env
+
+# Ensure uv is available (for local builds where base image lacks it)
+RUN if ! command -v uv >/dev/null 2>&1; then \
+        curl -LsSf https://astral.sh/uv/install.sh | sh && \
+        mv /root/.local/bin/uv /usr/local/bin/uv && \
+        mv /root/.local/bin/uvx /usr/local/bin/uvx; \
+    fi
+    
+# Install dependencies using uv sync
+# If uv.lock exists, use it; otherwise resolve on the fly
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-install-project --no-editable; \
+    else \
+        uv sync --no-install-project --no-editable; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-editable; \
+    else \
+        uv sync --no-editable; \
+    fi
+
+# Final runtime stage
+FROM ${BASE_IMAGE}
+
+WORKDIR /app
+
+# Copy the virtual environment from builder
+COPY --from=builder /app/env/.venv /app/.venv
+
+# Copy the environment code
+COPY --from=builder /app/env /app/env
+
+# Set PATH to use the virtual environment
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Set PYTHONPATH so imports work correctly
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+# Run the FastAPI server
+# The module path is constructed to work with the /app/env structure
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

README.md

@@ -0,0 +1,255 @@
+---
+title: Data Wrangler Environment Server
+emoji: 🎹
+colorFrom: green
+colorTo: yellow
+sdk: docker
+pinned: false
+app_port: 8000
+base_path: /web
+tags:
+  - openenv
+---
+
+# Data Wrangler Environment
+
+A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
+
+## Quick Start
+
+The simplest way to use the Data Wrangler environment is through the `DataWranglerEnv` class:
+
+```python
+from data_wrangler import DataWranglerAction, DataWranglerEnv
+
+try:
+    # Create environment from Docker image
+    data_wranglerenv = DataWranglerEnv.from_docker_image("data_wrangler-env:latest")
+
+    # Reset
+    result = data_wranglerenv.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+
+    # Send multiple messages
+    messages = ["Hello, World!", "Testing echo", "Final message"]
+
+    for msg in messages:
+        result = data_wranglerenv.step(DataWranglerAction(message=msg))
+        print(f"Sent: '{msg}'")
+        print(f"  → Echoed: '{result.observation.echoed_message}'")
+        print(f"  → Length: {result.observation.message_length}")
+        print(f"  → Reward: {result.reward}")
+
+finally:
+    # Always clean up
+    data_wranglerenv.close()
+```
+
+That's it! The `DataWranglerEnv.from_docker_image()` method handles:
+- Starting the Docker container
+- Waiting for the server to be ready
+- Connecting to the environment
+- Container cleanup when you call `close()`
+
+## Building the Docker Image
+
+Before using the environment, you need to build the Docker image:
+
+```bash
+# From project root
+docker build -t data_wrangler-env:latest -f server/Dockerfile .
+```
+
+## Deploying to Hugging Face Spaces
+
+You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
+
+```bash
+# From the environment directory (where openenv.yaml is located)
+openenv push
+
+# Or specify options
+openenv push --namespace my-org --private
+```
+
+The `openenv push` command will:
+1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
+2. Prepare a custom build for Hugging Face Docker space (enables web interface)
+3. Upload to Hugging Face (ensuring you're logged in)
+
+### Prerequisites
+
+- Authenticate with Hugging Face: The command will prompt for login if not already authenticated
+
+### Options
+
+- `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory)
+- `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
+- `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
+- `--private`: Deploy the space as private (default: public)
+
+### Examples
+
+```bash
+# Push to your personal namespace (defaults to username/env-name from openenv.yaml)
+openenv push
+
+# Push to a specific repository
+openenv push --repo-id my-org/my-env
+
+# Push with a custom base image
+openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest
+
+# Push as a private space
+openenv push --private
+
+# Combine options
+openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
+```
+
+After deployment, your space will be available at:
+`https://huggingface.co/spaces/<repo-id>`
+
+The deployed space includes:
+- **Web Interface** at `/web` - Interactive UI for exploring the environment
+- **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
+- **Health Check** at `/health` - Container health monitoring
+- **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
+
+## Environment Details
+
+### Action
+**DataWranglerAction**: Contains a single field
+- `message` (str) - The message to echo back
+
+### Observation
+**DataWranglerObservation**: Contains the echo response and metadata
+- `echoed_message` (str) - The message echoed back
+- `message_length` (int) - Length of the message
+- `reward` (float) - Reward based on message length (length × 0.1)
+- `done` (bool) - Always False for echo environment
+- `metadata` (dict) - Additional info like step count
+
+### Reward
+The reward is calculated as: `message_length × 0.1`
+- "Hi" → reward: 0.2
+- "Hello, World!" → reward: 1.3
+- Empty message → reward: 0.0
+
+## Advanced Usage
+
+### Connecting to an Existing Server
+
+If you already have a Data Wrangler environment server running, you can connect directly:
+
+```python
+from data_wrangler import DataWranglerEnv
+
+# Connect to existing server
+data_wranglerenv = DataWranglerEnv(base_url="<ENV_HTTP_URL_HERE>")
+
+# Use as normal
+result = data_wranglerenv.reset()
+result = data_wranglerenv.step(DataWranglerAction(message="Hello!"))
+```
+
+Note: When connecting to an existing server, `data_wranglerenv.close()` will NOT stop the server.
+
+### Using the Context Manager
+
+The client supports context manager usage for automatic connection management:
+
+```python
+from data_wrangler import DataWranglerAction, DataWranglerEnv
+
+# Connect with context manager (auto-connects and closes)
+with DataWranglerEnv(base_url="http://localhost:8000") as env:
+    result = env.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+    # Multiple steps with low latency
+    for msg in ["Hello", "World", "!"]:
+        result = env.step(DataWranglerAction(message=msg))
+        print(f"Echoed: {result.observation.echoed_message}")
+```
+
+The client uses WebSocket connections for:
+- **Lower latency**: No HTTP connection overhead per request
+- **Persistent session**: Server maintains your environment state
+- **Efficient for episodes**: Better for many sequential steps
+
+### Concurrent WebSocket Sessions
+
+The server supports multiple concurrent WebSocket connections. To enable this,
+modify `server/app.py` to use factory mode:
+
+```python
+# In server/app.py - use factory mode for concurrent sessions
+app = create_app(
+    DataWranglerEnvironment,  # Pass class, not instance
+    DataWranglerAction,
+    DataWranglerObservation,
+    max_concurrent_envs=4,  # Allow 4 concurrent sessions
+)
+```
+
+Then multiple clients can connect simultaneously:
+
+```python
+from data_wrangler import DataWranglerAction, DataWranglerEnv
+from concurrent.futures import ThreadPoolExecutor
+
+def run_episode(client_id: int):
+    with DataWranglerEnv(base_url="http://localhost:8000") as env:
+        result = env.reset()
+        for i in range(10):
+            result = env.step(DataWranglerAction(message=f"Client {client_id}, step {i}"))
+        return client_id, result.observation.message_length
+
+# Run 4 episodes concurrently
+with ThreadPoolExecutor(max_workers=4) as executor:
+    results = list(executor.map(run_episode, range(4)))
+```
+
+## Development & Testing
+
+### Direct Environment Testing
+
+Test the environment logic directly without starting the HTTP server:
+
+```bash
+# From the server directory
+python3 server/data_wrangler_environment.py
+```
+
+This verifies that:
+- Environment resets correctly
+- Step executes actions properly
+- State tracking works
+- Rewards are calculated correctly
+
+### Running Locally
+
+Run the server locally for development:
+
+```bash
+uvicorn server.app:app --reload
+```
+
+## Project Structure
+
+```
+data_wrangler/
+├── .dockerignore         # Docker build exclusions
+├── __init__.py            # Module exports
+├── README.md              # This file
+├── openenv.yaml           # OpenEnv manifest
+├── pyproject.toml         # Project metadata and dependencies
+├── uv.lock                # Locked dependencies (generated)
+├── client.py              # DataWranglerEnv client
+├── models.py              # Action and Observation models
+└── server/
+    ├── __init__.py        # Server module exports
+    ├── data_wrangler_environment.py  # Core environment logic
+    ├── app.py             # FastAPI application (HTTP + WebSocket endpoints)
+    └── Dockerfile         # Container image definition
+```

__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Data Wrangler Environment."""
+
+from .client import DataWranglerEnv
+from .models import DataWranglerAction, DataWranglerObservation
+
+__all__ = [
+    "DataWranglerAction",
+    "DataWranglerObservation",
+    "DataWranglerEnv",
+]

client.py

@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Data Wrangler Environment Client."""
+
+from typing import Dict
+
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+
+from .models import DataWranglerAction, DataWranglerObservation
+
+
+class DataWranglerEnv(
+    EnvClient[DataWranglerAction, DataWranglerObservation, State]
+):
+    """
+    Client for the Data Wrangler Environment.
+
+    This client maintains a persistent WebSocket connection to the environment server,
+    enabling efficient multi-step interactions with lower latency.
+    Each client instance has its own dedicated environment session on the server.
+
+    Example:
+        >>> # Connect to a running server
+        >>> with DataWranglerEnv(base_url="http://localhost:8000") as client:
+        ...     result = client.reset()
+        ...     print(result.observation.echoed_message)
+        ...
+        ...     result = client.step(DataWranglerAction(message="Hello!"))
+        ...     print(result.observation.echoed_message)
+
+    Example with Docker:
+        >>> # Automatically start container and connect
+        >>> client = DataWranglerEnv.from_docker_image("data_wrangler-env:latest")
+        >>> try:
+        ...     result = client.reset()
+        ...     result = client.step(DataWranglerAction(message="Test"))
+        ... finally:
+        ...     client.close()
+    """
+
+    def _step_payload(self, action: DataWranglerAction) -> Dict:
+        """
+        Convert DataWranglerAction to JSON payload for step message.
+
+        Args:
+            action: DataWranglerAction instance
+
+        Returns:
+            Dictionary representation suitable for JSON encoding
+        """
+        return {
+            "message": action.message,
+        }
+
+    def _parse_result(self, payload: Dict) -> StepResult[DataWranglerObservation]:
+        """
+        Parse server response into StepResult[DataWranglerObservation].
+
+        Args:
+            payload: JSON response data from server
+
+        Returns:
+            StepResult with DataWranglerObservation
+        """
+        obs_data = payload.get("observation", {})
+        observation = DataWranglerObservation(
+            echoed_message=obs_data.get("echoed_message", ""),
+            message_length=obs_data.get("message_length", 0),
+            done=payload.get("done", False),
+            reward=payload.get("reward"),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> State:
+        """
+        Parse server response into State object.
+
+        Args:
+            payload: JSON response from state request
+
+        Returns:
+            State object with episode_id and step_count
+        """
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )

inference.py

@@ -0,0 +1,190 @@
+import os
+import sys
+import asyncio
+from openai import AsyncOpenAI
+
+# OpenEnv V5 specific client components
+# We import directly since OpenEnv varies slightly in versions, but this mirrors the validator script expectations.
+try:
+    from openenv.core.client import EnvClient
+except ImportError:
+    pass
+
+API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
+API_KEY = os.environ.get("OPENAI_API_KEY", "")
+MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-3.5-turbo")
+IMAGE_NAME = "data_wrangler"
+TASK_NAME = "Data Writer Level 1"
+BENCHMARK = "data_wrangler"
+MAX_STEPS = 10
+MAX_TOTAL_REWARD = 1.0
+SUCCESS_SCORE_THRESHOLD = 0.5
+
+system_prompt = """\
+SYSTEM INSTRUCTIONS: ELITE DATA ENGINEER AGENT
+
+ROLE AND PERSONA
+You are an elite Data Engineering AI Agent operating within an automated data-wrangling pipeline. Your core function is to autonomously clean, format, and standardize messy, real-world datasets until they perfectly match a hidden "ground truth" target. You operate systematically, analytically, and with absolute precision.
+
+MISSION OBJECTIVE
+At each step, you will receive an Observation of the current data state. You must analyze the data anomalies (missing values, bad schemas, incorrect data types) and issue exactly ONE valid operation from your Action Space. You will iterate on this process until the dataset is perfectly clean, at which point you will issue the submit action.
+
+THE OBSERVATION
+You will receive a state dictionary detailing the dataset's current form:
+columns: Current list of headers.
+row_count: Total number of rows in the dataset.
+column_stats: Dictionary mapping column names to {dtype, missing_count, sample_values}.
+last_action_feedback: Status/error message resulting from your previous action.
+is_done: Boolean termination flag.
+
+ACTION SPACE (AVAILABLE TOOLS)
+You have a strict, highly constrained toolset. Your chosen action MUST be a valid JSON object matching exactly ONE of the schemas:
+1. Drop Column: {"action_type": "drop_column", "target_column": "..."}
+2. Rename Column: {"action_type": "rename_column", "target_column": "...", "new_name": "..."}
+3. Fill Missing Values: {"action_type": "fill_missing", "target_column": "...", "fill_value": "..."}
+4. Cast Data Type: {"action_type": "cast_type", "target_column": "...", "cast_to": "..."}
+5. Submit: {"action_type": "submit"}
+
+REQUIRED OUTPUT FORMAT (CHAIN OF THOUGHT)
+<thinking>
+Analyze Observation: What is the current state? What did the last action do?
+Identify Anomalies: Which columns have wrong types, bad names, or missing data?
+Formulate Plan: What is the highest priority fix right now?
+Select Action: Which action type and parameters will execute this fix?
+</thinking>
+{
+"action_type": "...",
+...
+}
+"""
+
+async def get_model_message(client, step, obs_dict, last_reward, history):
+    obs_text = str(obs_dict)
+    prompt = f"Step {step}.\nObservation: {obs_text}\nLast Reward: {last_reward}\nHistory: {history}\nChoose your next action (JSON matching schema)."
+    try:
+        response = await client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.0
+        )
+        content = response.choices[0].message.content
+        import json
+        import re
+        # Basic parsing of the JSON structure that follows the thinking tags
+        match = re.search(r'(\{.*\})', content, re.DOTALL)
+        if match:
+            return json.loads(match.group(1))
+        # Fallback if unparseable
+        return {"action_type": "submit"}
+    except Exception as e:
+        return {"action_type": "submit"}
+
+def log_start(task, env, model):
+    print(f"[START] task={task} env={env} model={model}")
+
+def log_step(step, action, reward, done, error):
+    print(f"[STEP] step={step} action={action} reward={reward} done={done} error={error}")
+
+def log_end(success, steps, score, rewards):
+    print(f"[END] success={success} steps={steps} score={score} rewards={rewards}")
+
+async def main():
+    if not API_KEY:
+        print("Missing OPENAI_API_KEY environment variable.")
+        return
+
+    client = AsyncOpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+    
+    print(f"[DEBUG] Spinning up {IMAGE_NAME} environment container...", flush=True)
+    try:
+        from client import DataWranglerEnv
+        env = DataWranglerEnv.from_docker_image(IMAGE_NAME)
+    except Exception as e:
+        print(f"[DEBUG] Docker env start failed ({e}). Falling back to local direct Python import.", flush=True)
+        from server.data_wrangler_environment import DataWranglerEnvironment
+        env = DataWranglerEnvironment() # Fallback for local debugging
+
+    history = []
+    rewards = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+
+    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
+
+    try:
+        if hasattr(env, 'reset') and not asyncio.iscoroutinefunction(env.reset):
+            result = env.reset()
+        else:
+            result = await env.reset() # OpenENV.reset() as per sample
+
+        obs = getattr(result, "observation", result)
+        obs_dict = {
+            "columns": getattr(obs, "columns", []),
+            "row_count": getattr(obs, "row_count", 0),
+            "column_stats": getattr(obs, "column_stats", {}),
+            "last_action_feedback": getattr(obs, "last_action_feedback", ""),
+            "is_done": getattr(obs, "is_done", False)
+        }
+        last_reward = getattr(result, "reward", getattr(obs, "reward", 0.0)) or 0.0
+
+        for step in range(1, MAX_STEPS + 1):
+            done = getattr(result, "done", getattr(obs, "is_done", False))
+            if done:
+                break
+
+            action_data = await get_model_message(client, step, obs_dict, last_reward, history)
+            
+            from models import DataWranglerAction
+            action_obj = DataWranglerAction(**action_data)
+            
+            if hasattr(env, 'step') and not asyncio.iscoroutinefunction(env.step):
+                result = env.step(action_obj)
+            else:
+                result = await env.step(action_obj)
+            
+            obs = getattr(result, "observation", result)
+            obs_dict = {
+                "columns": getattr(obs, "columns", []),
+                "row_count": getattr(obs, "row_count", 0),
+                "column_stats": getattr(obs, "column_stats", {}),
+                "last_action_feedback": getattr(obs, "last_action_feedback", ""),
+                "is_done": getattr(obs, "is_done", False)
+            }
+
+            reward = getattr(result, "reward", getattr(obs, "reward", 0.0)) or 0.0
+            done = getattr(result, "done", getattr(obs, "is_done", False))
+            error = None
+
+            rewards.append(reward)
+            steps_taken = step
+            last_reward = reward
+
+            log_step(step=step, action=action_data, reward=reward, done=done, error=error)
+
+            history.append(f"Step {step}: {action_data} -> reward {reward:+.2f}")
+
+            if done:
+                break
+
+        score = sum(rewards) / MAX_TOTAL_REWARD if MAX_TOTAL_REWARD > 0 else 0.0
+        score = min(max(score, 0.0), 1.0)
+        success = score >= SUCCESS_SCORE_THRESHOLD
+
+    finally:
+        try:
+            if hasattr(env, 'close'):
+                if asyncio.iscoroutinefunction(env.close):
+                    await env.close()
+                else:
+                    env.close()
+        except Exception as e:
+            print(f"[DEBUG] env.close() error (container cleanup): {e}", flush=True)
+            
+        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+
+if __name__ == "__main__":
+    asyncio.run(main())

models.py

@@ -0,0 +1,33 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Data models for the Data Wrangler Environment.
+
+The data_wrangler environment is a simple test environment that echoes back messages.
+"""
+
+from typing import Dict, List, Optional, Any
+from openenv.core.env_server.types import Action, Observation
+from pydantic import Field
+
+class DataWranglerAction(Action):
+    """Action for the Data Wrangler environment."""
+    action_type: str = Field(..., description="Type of action: drop_column, rename_column, fill_missing, cast_type, submit")
+    
+    # Specifics depending on action_type
+    target_column: Optional[str] = Field(None, description="The name of the column to act upon.")
+    new_name: Optional[str] = Field(None, description="New name of the column (for rename_column).")
+    fill_value: Optional[str] = Field(None, description="Value to fill missing data with (for fill_missing).")
+    cast_to: Optional[str] = Field(None, description="Target data type (for cast_type, e.g. 'int', 'float', 'datetime', 'string').")
+
+class DataWranglerObservation(Observation):
+    """Observation representing the state of the dataset."""
+    columns: List[str] = Field(default_factory=list, description="Current list of headers.")
+    row_count: int = Field(default=0, description="Total number of rows in the dataset.")
+    column_stats: Dict[str, Dict[str, Any]] = Field(default_factory=dict, description="Stats per column: dtype, missing_count, sample_values.")
+    last_action_feedback: str = Field(default="Environment initialized.", description="Feedback from the last executed action.")
+    is_done: bool = Field(default=False, description="Whether the task has terminated.")

openenv.yaml

@@ -0,0 +1,7 @@
+spec_version: 1
+name: data_wrangler
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000
+

pyproject.toml

@@ -0,0 +1,37 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-data_wrangler"
+version = "0.1.0"
+description = "Data Wrangler environment for OpenEnv"
+requires-python = ">=3.10"
+dependencies = [
+    # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+    "openenv-core[core]>=0.2.2",
+    "pandas>=2.0.0",
+    "openai>=1.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+# Server entry point - enables running via: uv run --project . server
+# or: python -m data_wrangler.server.app
+server = "data_wrangler.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["data_wrangler", "data_wrangler.server"]
+package-dir = { "data_wrangler" = ".", "data_wrangler.server" = "server" }

server/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Data Wrangler environment server components."""
+
+from .data_wrangler_environment import DataWranglerEnvironment
+
+__all__ = ["DataWranglerEnvironment"]

server/app.py

@@ -0,0 +1,91 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+FastAPI application for the Data Wrangler Environment.
+
+This module creates an HTTP server that exposes the DataWranglerEnvironment
+over HTTP and WebSocket endpoints, compatible with EnvClient.
+
+Endpoints:
+    - POST /reset: Reset the environment
+    - POST /step: Execute an action
+    - GET /state: Get current environment state
+    - GET /schema: Get action/observation schemas
+    - WS /ws: WebSocket endpoint for persistent sessions
+
+Usage:
+    # Development (with auto-reload):
+    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
+
+    # Production:
+    uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
+
+    # Or run directly:
+    python -m server.app
+"""
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:  # pragma: no cover
+    raise ImportError(
+        "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
+    ) from e
+
+try:
+    from ..models import DataWranglerAction, DataWranglerObservation
+    from .data_wrangler_environment import DataWranglerEnvironment
+except (ImportError, ValueError, ModuleNotFoundError):
+    import sys
+    import os
+    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+    from models import DataWranglerAction, DataWranglerObservation
+    from server.data_wrangler_environment import DataWranglerEnvironment
+
+
+# Create the app with web interface and README integration
+app = create_app(
+    DataWranglerEnvironment,
+    DataWranglerAction,
+    DataWranglerObservation,
+    env_name="data_wrangler",
+    max_concurrent_envs=1,  # increase this number to allow more concurrent WebSocket sessions
+)
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    """
+    Entry point for direct execution via uv run or python -m.
+
+    This function enables running the server without Docker:
+        uv run --project . server
+        uv run --project . server --port 8001
+        python -m data_wrangler.server.app
+
+    Args:
+        host: Host address to bind to (default: "0.0.0.0")
+        port: Port number to listen on (default: 8000)
+
+    For production deployments, consider using uvicorn directly with
+    multiple workers:
+        uvicorn data_wrangler.server.app:app --workers 4
+    """
+    import uvicorn
+
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == '__main__':
+    # ensures openenv validate passes main() check
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8000)
+    args = parser.parse_args()
+    if not args.port or args.port == 8000:
+        main()
+    else:
+        main(port=args.port)

server/data_wrangler_environment.py

@@ -0,0 +1,178 @@
+import os
+import pandas as pd
+from uuid import uuid4
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+try:
+    from ..models import DataWranglerAction, DataWranglerObservation
+except (ImportError, ValueError, ModuleNotFoundError):
+    import sys
+    import os
+    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+    from models import DataWranglerAction, DataWranglerObservation
+
+class DataWranglerEnvironment(Environment):
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(self):
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._reset_count = 0
+        self.df = None
+        self.target_df = None
+        self.task_level = int(os.environ.get("TASK_LEVEL", "1"))
+        self._initialize_task()
+
+    def _initialize_task(self):
+        self.df = pd.DataFrame()
+        self.target_df = pd.DataFrame()
+        if self.task_level == 1:
+            # Easy: Just drop a column and rename one
+            self.df = pd.DataFrame({
+                "User Name": ["Alice", "Bob", "Charlie"],
+                "Unnamed: 0": [0, 1, 2],
+                "Age": [25, 30, 35]
+            })
+            self.target_df = pd.DataFrame({
+                "user_name": ["Alice", "Bob", "Charlie"],
+                "age": [25, 30, 35]
+            })
+        elif self.task_level == 2:
+            # Medium: fill missing and cast type
+            self.df = pd.DataFrame({
+                "product_ID ": ["101", "102", "103"],
+                "price": ["10.5", None, "12.0"],
+                "bad_col": [None, None, None]
+            })
+            self.target_df = pd.DataFrame({
+                "product_id": [101.0, 102.0, 103.0], 
+                "price": [10.5, 0.0, 12.0]
+            })
+        else:
+            # Hard: Multiple issues
+            self.df = pd.DataFrame({
+                "date_joined ": ["2020-01-01", "2021-05-15", None],
+                "Sales_total": ["100", "200", "300"],
+                "IsActive": [True, False, None],
+                "DROPME_1": [1,2,3]
+            })
+            self.target_df = pd.DataFrame({
+                "date_joined": [pd.Timestamp("2020-01-01"), pd.Timestamp("2021-05-15"), pd.Timestamp("1970-01-01")],
+                "sales_total": [100.0, 200.0, 300.0],
+                "is_active": [True, False, False]
+            })
+
+    def _get_obs(self, feedback: str = "Environment initialized.", done: bool = False, reward: float = 0.0) -> DataWranglerObservation:
+        stats = {}
+        for col in self.df.columns:
+            stats[col] = {
+                "dtype": str(self.df[col].dtype),
+                "missing_count": int(self.df[col].isna().sum()),
+                "sample_values": self.df[col].dropna().astype(str).tolist()[:3]
+            }
+        
+        return DataWranglerObservation(
+            columns=list(self.df.columns),
+            row_count=len(self.df),
+            column_stats=stats,
+            last_action_feedback=feedback,
+            is_done=done,
+            reward=reward,
+            done=done,
+            metadata={"step": self._state.step_count}
+        )
+
+    def reset(self) -> DataWranglerObservation:
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._reset_count += 1
+        self._initialize_task()
+        return self._get_obs()
+
+    def step(self, action: DataWranglerAction) -> DataWranglerObservation: # type: ignore
+        self._state.step_count += 1
+        feedback = "Action executed successfully."
+        reward = 0.0
+        done = False
+        
+        try:
+            if action.action_type == "drop_column":
+                col = action.target_column
+                if col in self.df.columns:
+                    self.df.drop(columns=[col], inplace=True)
+                    if col not in self.target_df.columns:
+                        reward = 0.2
+                    else:
+                        reward = -0.5
+                        feedback = f"Warning: dropped targeting column {col}"
+                else:
+                    feedback = f"Error: Column '{col}' not found."
+            
+            elif action.action_type == "rename_column":
+                col = action.target_column
+                new_col = action.new_name
+                if col in self.df.columns:
+                    self.df.rename(columns={col: new_col}, inplace=True)
+                    if new_col in self.target_df.columns:
+                        reward = 0.2
+                else:
+                    feedback = f"Error: Column '{col}' not found."
+                    
+            elif action.action_type == "fill_missing":
+                col = action.target_column
+                if col in self.df.columns:
+                    self.df[col].fillna(action.fill_value, inplace=True)
+                    reward = 0.1
+                else:
+                    feedback = f"Error: Column '{col}' not found."
+                    
+            elif action.action_type == "cast_type":
+                col = action.target_column
+                to_type = action.cast_to
+                if col in self.df.columns:
+                    if to_type == 'int':
+                        self.df = self.df.astype({col: int})
+                    elif to_type == 'float':
+                        self.df = self.df.astype({col: float})
+                    elif to_type == 'datetime':
+                        self.df[col] = pd.to_datetime(self.df[col])
+                    elif to_type == 'string':
+                        self.df = self.df.astype({col: str})
+                    reward = 0.2
+                else:
+                    feedback = f"Error: Column '{col}' not found."
+                    
+            elif action.action_type == "submit":
+                score = self._grade()
+                reward = score
+                feedback = f"Submitted. Final Score: {score}"
+                done = True
+            else:
+                feedback = f"Error: Unknown action type {action.action_type}"
+                
+        except Exception as e:
+            feedback = f"Exception occurred: {str(e)}"
+            reward = -0.1
+
+        return self._get_obs(feedback=feedback, done=done, reward=reward)
+
+    def _grade(self) -> float:
+        score = 0.0
+        if list(self.df.columns) == list(self.target_df.columns):
+            score += 0.5
+             # Match types and values
+            value_matches = 0
+            for col in self.df.columns:
+                try:
+                    # simple match check
+                    match = (self.df[col] == self.target_df[col]).all()
+                    if match:
+                        value_matches += 1
+                except:
+                    pass
+            score += 0.5 * (value_matches / max(len(self.target_df.columns), 1))
+            
+        return score
+
+    @property
+    def state(self) -> State:
+        return self._state

server/requirements.txt

@@ -0,0 +1,8 @@
+openenv[core]>=0.2.0
+fastapi>=0.115.0
+uvicorn>=0.24.0
+
+
+
+pandas>=2.0.0
+openai>=1.0.0