Apppublic

AmbekarKush/etl-review-env

sourceHugging Facemitupdated 6mo agoView on Hugging Face

0likes

App README

🔧 Pipeline Doctor — AI ETL Debugging Environment

An AI agent acts as a data engineer on-call. It receives broken multi-stage ETL pipelines, diagnoses bugs, and applies targeted fixes — earning rewards proportional to output correctness.

OpenEnv environment | 4 tasks (easy → expert) | Hybrid pandas + SQLite | 41 tests

🌍 Environment Overview

Pipeline Doctor simulates a real-world scenario every data team faces: a production ETL pipeline has failed. The AI agent must inspect intermediate data, form hypotheses, and apply precise fixes — exactly like a senior data engineer debugging a failing Airflow DAG.

Why this domain is hard:

—Most failures are silent — no exception, just wrong data downstream
—Bugs cascade across stages: fixing one reveals the next
—Schema drift, type coercion, join key mismatches — all realistic failure modes

🧩 OpenEnv Compliance

Component	Status
`reset()` → initial observation	✅
`step(action)` → obs, reward, done, info	✅
`state()` → episode metadata	✅
Typed Pydantic models	✅
`openenv.yaml`	✅
`/tasks`, `/grader`, `/baseline` endpoints	✅
Dockerfile + HF Spaces	✅

🎯 Observation Space

Field	Type	Description
`pipeline_stages`	`List[StageDefinition]`	All stages with current code
`stage_results`	`List[StageResult]`	Runtime status: PASSING / FAILING / SKIPPED
`stage_statuses`	`Dict[str, str]`	Quick stage_id → status lookup
`current_score`	`float [0,1]`	Score from most recent fix
`source_schema`	`Dict[str, str]`	Column → dtype of raw input
`reasoning_trace`	`List[ReasoningTrace]`	Agent hypothesis + evidence per step
`efficiency_indicator`	`str`	`optimal` / `good` / `slow`
`fix_history`	`List[FixAttempt]`	All prior actions
`done`	`bool`	Episode terminated

⚡ Action Space

Action	Required	Description
`submit_fix`	`stage_id`, `fixed_code`	Apply corrected code, re-run pipeline, compute reward
`inspect_data`	`stage_id`	Request sample rows from stage output
`add_comment`	`comment`	Record diagnosis — earns up to +0.05 bonus
`skip`	—	Abandon episode

Pandas stage: Python block where df = input, assign result SQL stage: SELECT statement where input table = data

🏆 Reward Function

reward = (row_accuracy × 0.50)
       + (stage_pass_rate × 0.30 + column_match × 0.20) × efficiency_multiplier
       + comment_bonus (up to +0.05)

—Efficiency multiplier: 1.0 at steps 1-5, decays to 0.85 at step 10
—Applies only to structural scores — correct output always scores ≥0.99
—Partial progress always rewarded: fixing 1 of 2 bugs → ~0.35 score

📋 Tasks

Task	Difficulty	Stages	Bugs	Bug Type
`dtype_cast_failure`	🟢 Easy	3	1	Comma in numeric string → NaN
`multi_stage_logic_bug`	🟡 Medium	4	2	Silent case mismatch + wrong join key
`schema_drift_silent_loss`	🔴 Hard	5	3	Upstream renamed columns across 5 stages
`multi_source_reconciliation`	💀 Expert	5	4	Wrong dedup + outer join + fillna(999) + SQL rename

📊 Baseline Scores

Task	Difficulty	Score
`dtype_cast_failure`	🟢 Easy	~0.85
`multi_stage_logic_bug`	🟡 Medium	~0.63
`schema_drift_silent_loss`	🔴 Hard	~0.38
`multi_source_reconciliation`	💀 Expert	~0.22
Average		~0.52

`OPENAI_API_KEY=sk-... OPENAI_MODEL=gpt-4o-mini python baseline/run_baseline.py`

🎮 2-Minute Demo Walkthrough

bash

# 1. Reset
POST /reset  {"task_id": "dtype_cast_failure"}
# stage cast_revenue=FAILING, others=SKIPPED

# 2. Inspect
POST /step  {"action_type": "inspect_data", "stage_id": "cast_revenue"}
# sample rows show revenue=NaN

# 3. Diagnose
POST /step  {"action_type": "add_comment",
             "comment": "Comma thousands separator breaks pd.to_numeric"}

# 4. Fix
POST /step  {"action_type": "submit_fix", "stage_id": "cast_revenue",
             "fixed_code": "result = df.copy()\nresult['revenue'] = pd.to_numeric(result['revenue'].str.replace(',',''), errors='coerce')\n"}
# reward=1.0, done=True

# 5. Grade
POST /grader  →  {"score": 1.0, "stage_pass_rate": 1.0, "output_row_score": 1.0}

🚀 Setup

bash

# Local
pip install -r requirements.txt
uvicorn server.app:app --port 7860 --reload
streamlit run demo/streamlit_app.py   # separate terminal

# Docker (FastAPI + Streamlit)
docker build -t pipeline-doctor .
docker run -p 7860:7860 -p 8501:8501 pipeline-doctor

# API: http://localhost:7860/docs
# UI:  http://localhost:8501

🗂 Structure

etl-review-env/
├── env/             # models.py · environment.py · pipeline.py
├── tasks/           # task_easy · task_medium · task_hard · task_expert
├── graders/         # deterministic Jaccard grader
├── server/          # FastAPI (7 endpoints)
├── demo/            # Streamlit UI
├── baseline/        # OpenAI inference script
├── tests/           # 41 tests, 0 warnings
└── docker/          # supervisord (FastAPI + Streamlit)

MIT License