Blame for README.md

---
title: SQL Quality Audit Agent Environment
emoji: 🗄️
colorFrom: green
colorTo: blue
sdk: docker
pinned: false
tags:
  - openenv
  - sql
  - data-quality
  - rl-training
  - agent-evaluation
---

<div align="center">

# 🗄️ SQL Quality Audit Agent Environment

**The first multi-step RL environment purpose-built for SQL data quality reasoning.**

*Train agents that think like senior data engineers — not SQL autocomplete engines.*

[![OpenEnv](https://img.shields.io/badge/OpenEnv-Compatible-22c55e?style=flat-square)](https://huggingface.co/spaces/Deadman922/sql-quality-env)
[![Live Space](https://img.shields.io/badge/HuggingFace-Live%20Demo-f59e0b?style=flat-square)](https://huggingface.co/spaces/Deadman922/sql-quality-env)
[![License: MIT](https://img.shields.io/badge/License-MIT-3b82f6?style=flat-square)](LICENSE)
[![Tasks](https://img.shields.io/badge/Tasks-5%20Graded-8b5cf6?style=flat-square)](#tasks)

</div>

---

## The Problem We Are Solving

Every data team faces the same invisible crisis: production databases that silently degrade. Columns drift from `REAL` to `TEXT`. Foreign keys become orphaned. Duplicate rows inflate join counts by 2×. Schema versioning splits the same business concept across incompatible tables.

Existing AI benchmarks evaluate SQL *syntax*. None evaluate SQL *reasoning under uncertainty* — the ability to investigate an unknown database, form a hypothesis, apply a targeted fix, and verify the result before committing.

**This environment changes that.**

---

## What Makes This Different

| Benchmark Approach | This Environment |
|---|---|
| Static query → grade output | Multi-step episode with dense reward shaping |
| Clean, predefined schema | Unknown database with planted defects |
| Syntax correctness check | Five-dimension trajectory grading |
| One attempt per task | Configurable attempt budget per task type |
| No process feedback | Per-step rewards for investigation quality |

The grading system rewards *how* an agent works, not just *what* it submits. An agent that systematically investigates before fixing scores higher than one that submits the correct SQL on the first try — because in production, the correct SQL on the first try is often luck.

---

## Architecture Overview

```
┌─────────────────────────────────────────────────────────────┐
│                    AGENT (inference.py)                      │
│  LLM → generate_sql() → action_type + SQL/answer            │
└──────────────────────┬──────────────────────────────────────┘
                       │  POST /reset  POST /step
┌──────────────────────▼──────────────────────────────────────┐
│              ENVIRONMENT (FastAPI server)                    │
│  SQLQualityEnvironment → SQLite in-memory database          │
│                                                              │
│  QUERY ──→ read-only SELECT on live db                      │
│  PROFILE ─→ column statistics (null rate, type dist)        │
│  FIX ────→ write SQL on isolated _fresh_test_db() copy      │
│  SUBMIT ──→ compute_episode_score() → SQLReward             │
└──────────────────────┬──────────────────────────────────────┘
                       │
┌──────────────────────▼──────────────────────────────────────┐
│                 GRADER (graders.py)                          │
│  D1 Diagnostic · D2 Hypothesis · D3 Fix · D4 Verify · D5 Gen│
│  Composite score in [0.001, 0.999] → [END] stdout line      │
└─────────────────────────────────────────────────────────────┘
```

---

## Observation Space

Every `reset()` and `step()` returns a `SQLObservation` containing everything the agent needs to reason about the task:

```json
{
  "task_id": "dedup_bug",
  "task_description": "Find and repair the duplication bug inflating join result volumes",
  "schema_info": [
    {
      "table_name": "order_items",
      "columns": ["id (INTEGER)", "order_id (INTEGER)", "product_name (TEXT)", "quantity (INTEGER)", "price (REAL)"],
      "row_count": 320
    }
  ],
  "sample_rows": {
    "order_items": [{"id": 1, "order_id": 0, "product_name": "Product 1", "quantity": 2, "price": 19.99}]
  },
  "episode_id": "bf37e2ee-63e1-434f-9ff3-e0728518a81c",
  "attempt_number": 1,
  "max_attempts": 5,
  "trajectory_history": [],
  "available_actions": ["query", "profile", "fix", "submit"],
  "investigation_hint": null
}
```

| Field | Type | Description |
|---|---|---|
| `task_id` | string | Active task identifier |
| `task_description` | string | Natural language description of the defect to find and fix |
| `schema_info` | list | Table names, typed column names, and row counts |
| `sample_rows` | dict | First 5 rows of each table as JSON records |
| `episode_id` | string | UUID identifying the current episode |
| `trajectory_history` | list | All previous actions and result previews in this episode |
| `available_actions` | list | Valid action types at this point |
| `investigation_hint` | string or null | Actionable feedback after each scored step; null on first step |

---

## Action Space

Four actions are available at every step, each with strict database access controls:

```
┌──────────┬──────────────────────────┬─────────────────────────────────────┐
│ Action   │ Allowed SQL              │ Behavior                            │
├──────────┼──────────────────────────┼─────────────────────────────────────┤
│ QUERY    │ SELECT, WITH, EXPLAIN,   │ Read-only. Runs on live database.   │
│          │ PRAGMA                   │ Returns result preview. No state    │
│          │                          │ change. Call as many times as needed│
├──────────┼──────────────────────────┼─────────────────────────────────────┤
│ PROFILE  │ None (column name only)  │ Returns null rate, distinct count,  │
│          │                          │ and type distribution for one       │
│          │                          │ column. Shortcut for inspection.    │
├──────────┼──────────────────────────┼─────────────────────────────────────┤
│ FIX      │ UPDATE, INSERT, DELETE,  │ Runs on isolated _fresh_test_db()   │
│          │ CREATE, ALTER            │ copy. Scored but does NOT affect    │
│          │                          │ subsequent QUERY results.           │
├──────────┼──────────────────────────┼─────────────────────────────────────┤
│ SUBMIT   │ None                     │ Closes the episode. Triggers full   │
│          │                          │ trajectory scoring. Returns final   │
│          │                          │ composite reward.                   │
└──────────┴──────────────────────────┴─────────────────────────────────────┘
```

> **Key design decision:** FIX runs on `_fresh_test_db()` — an isolated in-memory copy that is discarded after scoring. The agent cannot verify its fix by re-querying the database afterward. It must reason about correctness from schema and sample data alone. This is precisely what makes D4 self-verification a non-trivial and RL-trainable behavior.

---

## The Five Tasks

### 🔵 `null_audit` — Easy — 1 attempt allowed

> *"One column in this table has more than 50% missing values. Find it."*

A `users` table with 8 columns where one randomly-selected column carries exactly a 52% null rate. All other columns have approximately 10% nulls as deliberate distractors. The agent must identify and name the high-null column. No SQL fix is required — this is the only pure-identification task in the benchmark.

**Why it is non-trivial:** Agents that profile columns one by one will exhaust their step budget. The efficient strategy is to query null rates for all columns in a single statement and immediately submit.

```sql
-- Efficient investigation
SELECT
  SUM(CASE WHEN email IS NULL THEN 1 ELSE 0 END) * 1.0 / COUNT(*) AS email_null_rate,
  SUM(CASE WHEN city IS NULL THEN 1 ELSE 0 END) * 1.0 / COUNT(*) AS city_null_rate
FROM users;
```

**Grader:** `server.graders:grade_null_audit`

---

### 🟡 `type_repair` — Medium — 3 attempts allowed

> *"Aggregate queries on this table return wrong totals. Something is wrong with the data types."*

A `transactions` table where the `amount` column is stored as `TEXT` containing numeric strings like `"12.50"`. SQLite aggregate functions silently return incorrect results because the column is treated as text. The fix must convert all 1,000 values to `REAL` without losing any rows. Data preservation is enforced at a 99.5% threshold.

**Why it is non-trivial:** SQLite does not support `ALTER COLUMN` syntax. Agents must use either `UPDATE ... SET amount = CAST(amount AS REAL)` or a full table recreation pattern. Many agents get stuck on invalid ALTER statements.

```sql
-- Correct approach
UPDATE transactions
SET amount = CAST(amount AS REAL)
WHERE typeof(amount) != 'real';
```

**Grader:** `server.graders:grade_type_repair`

---

### 🔴 `dedup_bug` — Hard — 5 attempts allowed

> *"A JOIN between orders and order_items is returning twice the expected row count. Track down why."*

An `order_items` table with exactly 320 rows — 160 legitimate records plus 160 exact duplicates, one per `order_id`. When joined against the `orders` table, every order appears twice. The correct clean count after deduplication is 160. The SUBMIT action accepts `root_cause_table` and `offending_key` fields for partial credit independent of fix quality.

**Why it is non-trivial:** The agent must correctly identify the duplication key (`order_id`), not just delete arbitrary rows. A naive `DELETE` without GROUP BY will produce incorrect results or violate the preservation threshold.

```sql
-- Correct approach
DELETE FROM order_items
WHERE rowid NOT IN (
    SELECT MIN(rowid) FROM order_items GROUP BY order_id
);
```

**Grader:** `server.graders:grade_dedup_bug`

---

### 🔴 `fk_integrity` — Hard — 4 attempts allowed

> *"Customer reports indicate missing invoice data. Referential integrity may be broken."*

An `invoices` table with 1,200 rows where 180 records reference `customer_id` values in the range 10,001–11,000 — IDs that do not exist in the `customers` table. Valid customer IDs are 1–400. These orphan records must be removed while preserving all 1,020 valid invoices intact. The fix is validated against both the orphan count and the preservation ratio.

```sql
-- Correct approach
DELETE FROM invoices
WHERE customer_id NOT IN (SELECT id FROM customers);
```

**Grader:** `server.graders:grade_fk_integrity`

---

### 🟡 `schema_drift` — Medium — 4 attempts allowed

> *"Two upstream data sources emit the same event concept but with completely different schemas. Build a canonical unified table."*

Two tables exist: `events_v1` (600 rows) with columns `[id, user_id, event_name, amount_text, event_time]` where `amount_text` is TEXT, and `events_v2` (600 rows) with columns `[event_id, account_id, event_type, amount, event_ts, source]` where `amount` is REAL. The agent must create a `canonical_events` table with 1,200 unified rows, correct column names, and consistent `REAL` type for the amount field.

**Why it is non-trivial:** Requires correctly mapping `id→event_id`, `user_id→user_id` (and `account_id→user_id` for events_v2), `event_name→event_type`, `CAST(amount_text AS REAL)→amount`, `event_time→event_time` across both sources. Incorrect column mapping produces wrong schema overlap scores.

```sql
-- Correct approach
CREATE TABLE canonical_events AS
    SELECT id AS event_id, user_id,
           event_name AS event_type,
           CAST(amount_text AS REAL) AS amount,
           event_time
    FROM events_v1
    UNION ALL
    SELECT event_id,
           account_id AS user_id,
           event_type,
           amount,
           event_ts AS event_time
    FROM events_v2;
```

**Grader:** `server.graders:grade_schema_drift`

---

## Five-Dimension Scoring

```
composite = D1 × 0.20  +  D2 × 0.20  +  D3 × 0.35  +  D4 × 0.10  +  D5 × 0.15
```

**D1 — Diagnostic Efficiency (20%)**
Measures how quickly the agent identified the defective table and column. Finding the issue in ≤2 queries scores 1.0. Three to four queries scores 0.7. Five to six queries scores 0.4. Never identifying the issue scores 0.0. This dimension directly penalizes shotgun exploration strategies.

**D2 — Hypothesis Quality (20%)**
Rewards structured investigation. Agents that inspect the table schema before drilling into specific columns score 1.0. Agents that query columns without understanding the table structure first score 0.6. This dimension trains broad-to-narrow reasoning habits.

**D3 — Fix Correctness (35%)**
The highest-weighted dimension. Evaluated with four severity tiers: PERFECT (issue fully resolved, all data preserved) = 1.0, ACCEPTABLE (issue resolved with minor data impact) = 0.6, DANGEROUS (partial fix that risks data loss) = 0.2, CATASTROPHIC (fix causes data destruction or no improvement) = 0.0. CATASTROPHIC severity overrides the entire composite and caps the final score at 0.05.

**D4 — Self-Verification (10%)**
A bonus dimension awarded when an agent submits a PROFILE action after a FIX to verify schemas and reason about correctness before calling SUBMIT. Because FIX runs on an isolated copy, the agent cannot re-query data, so it must perform reasoning-only checks. This behavior does not emerge from supervised learning on static SQL datasets. It is the primary RL-emergent behavior this environment is designed to cultivate. All zero-shot baseline agents score 0.0 here.

**D5 — Generalization (15%)**
The fix is silently tested against three schema variants: (1) a base copy of the original data, (2) the same schema with an additional column added, and (3) the same schema with a shifted data distribution. A fix that is brittle to schema changes scores 0.0 even if it perfectly resolves the original task.

### CATASTROPHIC Override

Any fix meeting one of the following conditions is immediately classified as CATASTROPHIC and caps the composite at **0.05** regardless of D1, D2, D4, and D5 scores:

- The fix contains `DROP TABLE` on the target table
- The fix results in zero rows remaining in the target table
- The fix deletes more than 50% of originally valid records

This prevents agents from gaming the reward signal through destructive operations.

### Score Boundaries

All scores are clamped to `[0.001, 0.999]`. Exact values of 0.0 and 1.0 are never returned. This design choice prevents RL training instability caused by absorbing boundary states.

---

## Setup and Installation

### Prerequisites

```bash
git clone https://huggingface.co/spaces/Deadman922/sql-quality-env
cd sql-quality-env
pip install -r requirements.txt
```

### Environment Variables

| Variable | Required | Default | Purpose |
|---|---|---|---|
| `API_BASE_URL` | Yes | `https://router.huggingface.co/v1` | LLM API endpoint |
| `MODEL_NAME` | Yes | `meta-llama/Llama-3.3-70B-Instruct` | Model identifier |
| `HF_TOKEN` | Yes | — | HuggingFace API token — no default, must be set explicitly |
| `LOCAL_IMAGE_NAME` | No | — | Docker image name, only required with `from_docker_image()` |

```bash
export API_BASE_URL="https://router.huggingface.co/v1"
export MODEL_NAME="meta-llama/Llama-3.3-70B-Instruct"
export HF_TOKEN="hf_your_token_here"
```

### Start the Server

```bash
# Python
python -m server.main

# Docker
docker build -t sql-quality-env .
docker run -p 7860:7860 \
  -e API_BASE_URL=$API_BASE_URL \
  -e MODEL_NAME=$MODEL_NAME \
  -e HF_TOKEN=$HF_TOKEN \
  sql-quality-env
```

### Run the Agent

```bash
# Single task
python inference.py --task null_audit --url http://localhost:7860

# Full benchmark run
for task in null_audit type_repair dedup_bug fk_integrity schema_drift; do
    echo "--- Running $task ---"
    python inference.py --task $task --url http://localhost:7860
done
```

### Verify the Environment

```bash
# Health check
curl http://localhost:7860/health

# List graded tasks
curl http://localhost:7860/tasks

# Start a manual episode
curl -X POST http://localhost:7860/reset \
  -H "Content-Type: application/json" \
  -d '{"task_id": "null_audit", "seed": 42}'
```

---

## Stdout Log Format

The inference agent emits exactly three structured line types. Any deviation in format causes evaluation failure at the validator:

```
[START] task=null_audit env=sql_quality_env model=meta-llama/Llama-3.3-70B-Instruct
[STEP]  step=1 action=query:SELECT COUNT(*)... reward=0.01 done=false error=null
[STEP]  step=2 action=submit:email reward=0.99 done=true error=null
[END]   success=true steps=2 score=0.999 rewards=0.01,0.99
```

| Field | Format | Notes |
|---|---|---|
| `reward` per step | 2 decimal places | `0.01` not `0.0105` |
| `score` in END | 3 decimal places | `0.999` not `1.0` |
| `done`, `success` | lowercase boolean | `true` or `false` never `True` |
| `error` | raw string or `null` | never `None` or empty string |
| debug output | stderr only | never appears in stdout |

---

## Baseline Scores

Evaluated with `meta-llama/Llama-3.3-70B-Instruct`, temperature 0.7, max 200 tokens, seed 42, force-submit at step 7:

| Task | Score | D1 | D2 | D3 | D4 | D5 | Primary Bottleneck |
|---|---|---|---|---|---|---|---|
| `null_audit` | ~0.68 | 1.00 | 0.60 | ~0.70 | 0.00 | 0.00 | Lexical near-misses on column name spelling |
| `type_repair` | ~0.41 | 1.00 | 0.60 | ~0.30 | 0.00 | 0.00 | SQLite ALTER limitation requires table recreation |
| `dedup_bug` | ~0.82 | 1.00 | 0.60 | 1.00 | 0.00 | 1.00 | Reliable once duplicate pattern is identified |
| `fk_integrity` | ~0.82 | 1.00 | 0.60 | 1.00 | 0.00 | 1.00 | Reliable with LEFT JOIN orphan detection |
| `schema_drift` | ~0.50 | 0.00 | 0.00 | 1.00 | 0.00 | 1.00 | No investigation steps before fix |

D4 (self-verification) scores 0.0 for all zero-shot agents. This is the benchmark's primary training signal. The behavior is RL-emergent and cannot be obtained through imitation learning.

**Reproducibility configuration:**
- Model: `meta-llama/Llama-3.3-70B-Instruct`
- Temperature: 0.7 · Max tokens: 200 · Top-p: 1.0
- Seed: 42 per task · Max steps: 20 · Force submit at step 7

---

## API Reference

| Endpoint | Method | Request Body | Response |
|---|---|---|---|
| `/health` | GET | — | `{"status": "ok", "llm_available": bool, "provider": string}` |
| `/tasks` | GET | — | `{"tasks": [...], "total": 5, "graded": 5}` |
| `/metadata` | GET | — | Environment metadata with full task registry |
| `/reset` | POST | `{"task_id": string, "seed"?: int}` | `SQLObservation` |
| `/step` | POST | `SQLAction` | `{"observation", "reward", "done", "info"}` |
| `/state` | GET | — | `{"task_id", "episode_id", "attempt_number", "cumulative_reward", "done"}` |

---

## Grader Registry

| Task | Grader Function | Method |
|---|---|---|
| `null_audit` | `server.graders:grade_null_audit` | String match with lexical partial credit |
| `type_repair` | `server.graders:grade_type_repair` | Schema + row preservation on isolated copy |
| `dedup_bug` | `server.graders:grade_dedup_bug` | Row count + key identification + fix validation |
| `fk_integrity` | `server.graders:grade_fk_integrity` | Orphan elimination + valid record preservation |
| `schema_drift` | `server.graders:grade_schema_drift` | Column overlap ratio + row count alignment |

---

<div align="center">

*Built for the Meta × PyTorch × HuggingFace OpenEnv Hackathon*

**One Man Automation · Chennai, India**

[Live Space](https://huggingface.co/spaces/Deadman922/sql-quality-env) · [Report Issue](https://huggingface.co/spaces/Deadman922/sql-quality-env/discussions)

</div>