README.md

metadata

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

🗄️ 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.

---

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:

{
  "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.

-- 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.

-- 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.

-- 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.

-- 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.

-- 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

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()

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

# 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

# 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

# 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

---

Built for the Meta × PyTorch × HuggingFace OpenEnv Hackathon

One Man Automation · Chennai, India

Live Space · Report Issue