naivaidhya/sql-data-quality-env

1---2title: SQL Data Quality Env3emoji: 🗄️4colorFrom: blue5colorTo: purple6sdk: docker7pinned: false8---9 10# 🗄️ SQL Data Quality Environment11 12<div align="center">13 14[![OpenEnv](https://img.shields.io/badge/OpenEnv-compatible-6c63ff)](https://github.com/meta-pytorch/OpenEnv)15[![HF Spaces](https://img.shields.io/badge/HuggingFace-Space-orange)](https://huggingface.co/spaces)16[![Python](https://img.shields.io/badge/Python-3.10+-blue)](https://python.org)17[![License](https://img.shields.io/badge/License-MIT-green)](LICENSE)18 19**A real-world OpenEnv environment where AI agents audit SQL databases for data quality issues and generate corrective SQL queries.**20 21</div>22 23---24 25## 🎯 What This Environment Simulates26 27Data quality issues cost enterprises an estimated **$12.9M per year** on average (Gartner). Every data engineer and analytics team spends significant time:28 29- Detecting NULL values in critical columns30- Finding and deduplicating records with the same business key31- Identifying referential integrity violations (orphan foreign keys)32- Fixing type mismatches and format inconsistencies33- Correcting calculation errors in derived columns34 35This environment reproduces exactly these real workflows through an interactive SQLite database that agents explore using `list_tables`, `describe_table`, `query`, and `submit_fix` actions — the same tools a real data engineer would use.36 37---38 39## 🏗️ Architecture40 41```42┌──────────────────────────────────────────────────────────┐43│                    Agent (LLM / RL)                      │44│          Observes text + structured JSON data            │45└────────────────────┬─────────────────────────────────────┘46                     │ HTTP  (reset / step / state)47┌────────────────────▼─────────────────────────────────────┐48│              FastAPI Server  (port 7860)                  │49│  ┌────────────────────────────────────────────────────┐  │50│  │         SQLDataQualityEnvironment                   │  │51│  │  reset() → fresh in-memory SQLite DB per episode   │  │52│  │  step()  → routes action, returns Observation      │  │53│  │  state   → episode metadata (step_count, score)    │  │54│  └────────────────────────────────────────────────────┘  │55│                                                          │56│  ┌──────────┐  ┌──────────┐  ┌──────────────────────┐   │57│  │ EASY     │  │ MEDIUM   │  │ HARD                 │   │58│  │1 table   │  │2 tables  │  │3 tables              │   │59│  │null/type │  │dedup+FK  │  │multi-table+biz rules │   │60│  └──────────┘  └──────────┘  └──────────────────────┘   │61└──────────────────────────────────────────────────────────┘62```63 64---65 66## 📐 OpenEnv Spec Compliance67 68| Requirement | Status |69|---|---|70| Typed `Action` Pydantic model | ✅ |71| Typed `Observation` Pydantic model | ✅ |72| Typed `State` Pydantic model | ✅ |73| `POST /reset` endpoint | ✅ |74| `POST /step` endpoint | ✅ |75| `GET /state` endpoint | ✅ |76| `openenv.yaml` manifest | ✅ |77| 3+ tasks with graders (0.0–1.0) | ✅ |78| Graders are deterministic | ✅ |79| Baseline inference script | ✅ |80| Working Dockerfile | ✅ |81| HF Spaces deployable | ✅ |82 83---84 85## 🎮 Action Space86 87The agent communicates through structured `Action` objects:88 89```python90class Action(BaseModel):91    action_type: ActionType        # required92    table_name:  Optional[str]     # for describe_table 93    sql:         Optional[str]     # for query (read-only SELECT)94    fix_sql:     Optional[str]     # for submit_fix (UPDATE/DELETE)95    reasoning:   Optional[str]     # optional chain-of-thought (not graded)96```97 98### Action Types99 100| `action_type` | Description | Required Fields |101|---|---|---|102| `list_tables` | List all tables in the database | — |103| `describe_table` | Schema + sample rows for a table | `table_name` |104| `query` | Execute a read-only SELECT/WITH | `sql` |105| `submit_fix` | Apply UPDATE/DELETE/ALTER fix statements | `fix_sql` |106| `finish` | End the episode, receive final score | — |107 108### Example Actions109 110```json111// Explore112{"action_type": "list_tables"}113{"action_type": "describe_table", "table_name": "customers"}114{"action_type": "query", "sql": "SELECT * FROM customers WHERE email IS NULL"}115 116// Fix117{"action_type": "submit_fix", "fix_sql": "UPDATE customers SET email='unknown@example.com' WHERE email IS NULL;"}118{"action_type": "finish"}119```120 121---122 123## 👁️ Observation Space124 125Each step returns an `Observation`:126 127```python128class Observation(BaseModel):129    done:             bool                   # True when episode ends130    reward:           float                  # Step reward (with shaping)131    observation_text: str                    # Human-readable description132    data: Optional[Dict[str, Any]]           # Structured payload (varies by action)133    error: Optional[str]                     # Error message if action failed134```135 136### Data Payload by Action137 138| Action | `data` structure |139|---|---|140| `list_tables` | `{tables: [str]}` |141| `describe_table` | `{columns: [{name, type, nullable}], sample_rows: [dict], total_rows: int}` |142| `query` | `{columns: [str], rows: [[...]], row_count: int}` |143| `submit_fix` | `{affected_rows: int, validation: {score_before, score_after, delta, errors}}` |144| `finish` | `{final_score: float, steps_used: int, efficiency_bonus: float, breakdown: dict}` |145 146---147 148## 📊 Episode State149 150```python151class State(BaseModel):152    episode_id:        str        # unique episode identifier153    task_id:           str        # 'easy' | 'medium' | 'hard'154    step_count:        int        # steps taken so far155    max_steps:         int        # budget (easy=20, medium=25, hard=35)156    cumulative_reward: float      # total reward accumulated157    issues_found:      int        # issues the agent has found158    fixes_applied:     int        # successful fix statements applied159    task_description:  str        # full task instructions160    available_tables:  List[str]  # tables in this episode's DB161```162 163---164 165## 📋 Tasks166 167### Task 1 — EASY: Customer Table Null & Type Audit168**Difficulty:** Easy | **Max steps:** 20 | **Table:** `customers`169 170The `customers` table has 10 rows with four categories of data quality issues:171 172| Category | Issue | Points |173|---|---|---|174| A | 3 rows missing `email` | 25% |175| B | 2 rows missing `phone` | 25% |176| C | 2 rows missing `city` or `country` | 25% |177| D | 3 rows with non-numeric `age` values (`'abc'`, `''`, `'N/A'`) | 25% |178 179**Expected agent strategy:** Describe table → query for NULLs by column → submit UPDATE fixes → finish180 181**Baseline score (GPT-4o-mini):** ~0.75182 183---184 185### Task 2 — MEDIUM: Products & Orders Integrity186**Difficulty:** Medium | **Max steps:** 25 | **Tables:** `products`, `orders`187 188Two related tables with three categories of issues:189 190| Category | Issue | Points |191|---|---|---|192| A | 2 duplicate SKUs (SKU-001, SKU-002) — keep lower product_id | 33% |193| B | 2 orphan orders referencing non-existent products | 33% |194| C | 2 orders with `quantity ≤ 0` (0 and -1) | 34% |195 196**Expected agent strategy:** Detect cross-table relationships → find duplicates with GROUP BY → identify FK violations → fix sequentially197 198**Baseline score (GPT-4o-mini):** ~0.60199 200---201 202### Task 3 — HARD: Multi-Table Schema & Business Rules203**Difficulty:** Hard | **Max steps:** 35 | **Tables:** `employees`, `departments`, `payroll`204 205Seven distinct issue categories across three tables:206 207| # | Category | Table | Issue |208|---|---|---|---|209| 1 | Date format | `employees` | Row 4: `hire_date='15/06/2018'` → ISO 8601 |210| 2 | Invalid FK | `employees` | Row 6: department `'Logistics'` not in `departments` |211| 3 | Self-ref integrity | `employees` | Row 7: `manager_id=99` doesn't exist |212| 4 | Business rule | `employees` | Row 8: negative `salary=-5000` |213| 5 | NULL constraint | `employees` | Row 11: NULL `department` |214| 6 | Calculation error | `payroll` | Rows 2, 5: `net_pay ≠ gross_pay - deductions` |215| 7 | Orphan record | `payroll` | Row 9: `emp_id=99` doesn't exist in `employees` |216 217**Expected agent strategy:** Deep multi-table analysis, verify self-referencing integrity, check derived column calculations, fix issues in dependency order218 219**Baseline score (GPT-4o-mini):** ~0.43220 221---222 223## 🏆 Reward Function224 225The reward function provides **dense, shaped signals** throughout the episode:226 227```228step_reward = -0.005                          # per-step efficiency penalty229            + 0.05                            # if fix affected ≥ 1 row(s)230            + max(grader_delta, 0.0) × 0.5   # proportional to quality improvement231```232 233On `finish()`:234```235finish_reward = final_grader_score            # 0.0 – 1.0236              + 0.10                          # efficiency bonus (≤ half step budget)237```238 239On timeout (step budget exhausted):240```241timeout_penalty = -0.05242```243 244**Design rationale:**245- The per-step penalty discourages aimless exploration without preventing necessary investigation246- `submit_fix` gives immediate feedback even before the episode ends247- The grader delta component rewards *meaningful* fixes, not just any SQL execution248- The efficiency bonus incentivises concise, targeted agents over brute-force approaches249 250---251 252## 🚀 Quick Start253 254### Local Setup255 256```bash257# Clone / download the project258cd scaler/259 260# Install dependencies261pip install -r requirements.txt262 263# Start the server264uvicorn server.app:app --host 0.0.0.0 --port 7860265 266# Open http://localhost:7860 in your browser for the web UI267# OpenAPI docs at http://localhost:7860/docs268```269 270### Docker271 272```bash273docker build -t sql-data-quality-env .274docker run -p 7860:7860 sql-data-quality-env275 276# Health check277curl http://localhost:7860/health278```279 280### Python Client281 282```python283from client import SQLDataQualityClient284from models import Action, ActionType285 286with SQLDataQualityClient("http://localhost:7860") as client:287    # Start easy task288    obs = client.reset(task_id="easy")289    print(obs.observation_text)290 291    # Explore292    result = client.step(Action(action_type=ActionType.LIST_TABLES))293    result = client.step(Action(action_type=ActionType.DESCRIBE_TABLE, table_name="customers"))294 295    # Fix296    result = client.step(Action(297        action_type=ActionType.SUBMIT_FIX,298        fix_sql="UPDATE customers SET email='unknown@example.com' WHERE email IS NULL;"299    ))300    print(result.observation.data["validation"])301 302    # Finish303    result = client.step(Action(action_type=ActionType.FINISH))304    print(f"Final score: {result.observation.data['final_score']}")305```306 307---308 309## 🤖 Baseline Inference Script310 311The baseline script runs a GPT-4o-mini agent through all three tasks:312 313```bash314# Required environment variables315export HF_TOKEN="your-api-key"316export API_BASE_URL="https://api.openai.com/v1"   # or your custom endpoint317export MODEL_NAME="gpt-4o-mini"                    # or your model318 319# Run all tasks (server must be running)320python inference.py --url http://localhost:7860321 322# Run specific tasks323python inference.py --tasks easy medium324 325# Quiet mode (just scores)326python inference.py --quiet327```328 329### Baseline Scores (Reproducible)330 331| Task | GPT-4o-mini | Notes |332|---|---|---|333| Easy | ~0.75 | Misses some edge cases in age validation |334| Medium | ~0.60 | Struggles with ordering of operations (must delete dupes before orphan check) |335| Hard | ~0.43 | Date format conversion and multi-table coordination are challenging |336| **Average** | **~0.59** | Room for significant improvement |337 338---339 340## 🧪 Running Tests341 342```bash343pytest tests/ -v344 345# Expected output:346# test_reset_easy_returns_observation PASSED347# test_reset_sets_state PASSED348# ...349# 22 passed in X.XX seconds350```351 352---353 354## 🐳 Deploying to Hugging Face Spaces355 3561. Create a new HF Space with **Docker** SDK3572. Push this repository to the Space3583. The Space will automatically build and start on port 78603594. Tag your Space with `openenv` for discoverability360 361The web UI at `/` provides a no-code interface for manual interaction.362 363---364 365## 📁 Project Structure366 367```368scaler/369├── openenv.yaml              # OpenEnv manifest370├── models.py                 # Pydantic models (Action, Observation, State, StepResult)371├── tasks.py                  # Task definitions, schemas, seed data, graders372├── environment.py            # Core environment logic (reset/step/state)373├── client.py                 # Synchronous HTTP client374├── inference.py              # Baseline inference script (OpenAI client)375├── requirements.txt          # Python dependencies376├── pyproject.toml            # Package metadata377├── Dockerfile                # Container definition378├── server/379│   └── app.py                # FastAPI application380└── tests/381    └── test_environment.py   # Test suite (pytest)382```383 384---385 386## 🔧 Environment Variables387 388| Variable | Description | Default |389|---|---|---|390| `API_BASE_URL` | LLM API endpoint | `https://api.openai.com/v1` |391| `MODEL_NAME` | Model identifier for inference | `gpt-4o-mini` |392| `HF_TOKEN` | API key (used as OpenAI `api_key`) | *(required)* |393 394---395 396## 🙏 Acknowledgments397 398Built for the [OpenEnv Community Challenge](https://github.com/meta-pytorch/OpenEnv) by Meta PyTorch × Hugging Face.399 400The SQL schema designs are inspired by real data quality issues encountered in production data warehouses.401