aarjavjjain/SQL-Data-Quality-OpenEnv

┌──────────────────────────────────────────────────────────┐
│                    Agent (LLM / RL)                      │
│          Observes text + structured JSON data            │
└────────────────────┬─────────────────────────────────────┘
                     │ HTTP  (reset / step / state)
┌────────────────────▼─────────────────────────────────────┐
│              FastAPI Server  (port 7860)                  │
│  ┌────────────────────────────────────────────────────┐  │
│  │         SQLDataQualityEnvironment                   │  │
│  │  reset() → fresh in-memory SQLite DB per episode   │  │
│  │  step()  → routes action, returns Observation      │  │
│  │  state   → episode metadata (step_count, score)    │  │
│  └────────────────────────────────────────────────────┘  │
│                                                          │
│  ┌──────────┐  ┌──────────┐  ┌──────────────────────┐   │
│  │ EASY     │  │ MEDIUM   │  │ HARD                 │   │
│  │1 table   │  │2 tables  │  │3 tables              │   │
│  │null/type │  │dedup+FK  │  │multi-table+biz rules │   │
│  └──────────┘  └──────────┘  └──────────────────────┘   │
└──────────────────────────────────────────────────────────┘

class Action(BaseModel):
    action_type: ActionType        # required
    table_name:  Optional[str]     # for describe_table 
    sql:         Optional[str]     # for query (read-only SELECT)
    fix_sql:     Optional[str]     # for submit_fix (UPDATE/DELETE)
    reasoning:   Optional[str]     # optional chain-of-thought (not graded)

// Explore
{"action_type": "list_tables"}
{"action_type": "describe_table", "table_name": "customers"}
{"action_type": "query", "sql": "SELECT * FROM customers WHERE email IS NULL"}

// Fix
{"action_type": "submit_fix", "fix_sql": "UPDATE customers SET email='unknown@example.com' WHERE email IS NULL;"}
{"action_type": "finish"}

class Observation(BaseModel):
    done:             bool                   # True when episode ends
    reward:           float                  # Step reward (with shaping)
    observation_text: str                    # Human-readable description
    data: Optional[Dict[str, Any]]           # Structured payload (varies by action)
    error: Optional[str]                     # Error message if action failed

class State(BaseModel):
    episode_id:        str        # unique episode identifier
    task_id:           str        # 'easy' | 'medium' | 'hard'
    step_count:        int        # steps taken so far
    max_steps:         int        # budget (easy=20, medium=25, hard=35)
    cumulative_reward: float      # total reward accumulated
    issues_found:      int        # issues the agent has found
    fixes_applied:     int        # successful fix statements applied
    task_description:  str        # full task instructions
    available_tables:  List[str]  # tables in this episode's DB

step_reward = -0.005                          # per-step efficiency penalty
            + 0.05                            # if fix affected ≥ 1 row(s)
            + max(grader_delta, 0.0) × 0.5   # proportional to quality improvement

finish_reward = final_grader_score            # 0.0 – 1.0
              + 0.10                          # efficiency bonus (≤ half step budget)

timeout_penalty = -0.05

# Clone / download the project
cd scaler/

# Install dependencies
pip install -r requirements.txt

# Start the server
uvicorn server.app:app --host 0.0.0.0 --port 7860

# Open http://localhost:7860 in your browser for the web UI
# OpenAPI docs at http://localhost:7860/docs

docker build -t sql-data-quality-env .
docker run -p 7860:7860 sql-data-quality-env

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

from client import SQLDataQualityClient
from models import Action, ActionType

with SQLDataQualityClient("http://localhost:7860") as client:
    # Start easy task
    obs = client.reset(task_id="easy")
    print(obs.observation_text)

    # Explore
    result = client.step(Action(action_type=ActionType.LIST_TABLES))
    result = client.step(Action(action_type=ActionType.DESCRIBE_TABLE, table_name="customers"))

    # Fix
    result = client.step(Action(
        action_type=ActionType.SUBMIT_FIX,
        fix_sql="UPDATE customers SET email='unknown@example.com' WHERE email IS NULL;"
    ))
    print(result.observation.data["validation"])

    # Finish
    result = client.step(Action(action_type=ActionType.FINISH))
    print(f"Final score: {result.observation.data['final_score']}")

# Required environment variables
export HF_TOKEN="your-api-key"
export API_BASE_URL="https://api.openai.com/v1"   # or your custom endpoint
export MODEL_NAME="gpt-4o-mini"                    # or your model

# Run all tasks (server must be running)
python inference.py --url http://localhost:7860

# Run specific tasks
python inference.py --tasks easy medium

# Quiet mode (just scores)
python inference.py --quiet

pytest tests/ -v

# Expected output:
# test_reset_easy_returns_observation PASSED
# test_reset_sets_state PASSED
# ...
# 22 passed in X.XX seconds

scaler/
├── openenv.yaml              # OpenEnv manifest
├── models.py                 # Pydantic models (Action, Observation, State, StepResult)
├── tasks.py                  # Task definitions, schemas, seed data, graders
├── environment.py            # Core environment logic (reset/step/state)
├── client.py                 # Synchronous HTTP client
├── inference.py              # Baseline inference script (OpenAI client)
├── requirements.txt          # Python dependencies
├── pyproject.toml            # Package metadata
├── Dockerfile                # Container definition
├── server/
│   └── app.py                # FastAPI application
└── tests/
    └── test_environment.py   # Test suite (pytest)