File size: 5,604 Bytes
1a1713a
e965a47
1a1713a
e965a47
f6bd747
e965a47
6c6f994
e965a47
 
 
 
 
 
 
 
 
 
 
f52aed3
e965a47
 
 
 
 
 
82a5b1b
e965a47
2194233
e965a47
f52aed3
 
2194233
 
 
f52aed3
7a6f18c
 
 
 
 
82a5b1b
7a6f18c
2194233
 
fd851a9
 
 
 
2194233
fd851a9
 
 
2194233
 
 
 
 
2be7532
2194233
 
82a5b1b
 
 
fd851a9
82a5b1b
fd851a9
2194233
82a5b1b
2194233
 
fd851a9
82a5b1b
2194233
 
fd851a9
 
 
 
f52aed3
 
 
 
 
 
 
 
 
82a5b1b
f52aed3
 
 
 
 
 
 
 
 
 
 
2194233
 
 
 
 
 
 
 
 
 
 
f52aed3
 
 
 
2194233
 
 
 
 
 
f52aed3
2194233
 
f52aed3
 
 
e965a47
f6bd747
e965a47
 
 
f6bd747
c4a66be
 
 
2194233
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1a1713a
 
 
 
 
 
66cc86e
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
"""
FastAPI server using openenv.core base classes β€” required for validator.
"""
import random
from openenv.core.env_server.http_server import create_app
from openenv.core.env_server.interfaces import Environment

try:
    from sql_env.models import SQLAction, SQLObservation, SQLState
    from sql_env.tasks import TASK_SETS
    from sql_env.grader import grade, generate_feedback
except ImportError:
    from models import SQLAction, SQLObservation, SQLState
    from tasks import TASK_SETS
    from grader import grade, generate_feedback


class SQLCorrectionEnvironment(Environment):
    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self):
        self._difficulty = "easy"
        self._current_task = None
        self._step_count = 0
        self._done = False
        self._last_reward = 0.01
        self._rewards_history = []
        self._stagnation_count = 0

    def reset(self, seed=None, episode_id=None, **kwargs) -> SQLObservation:
        actual_difficulty = (
            kwargs.get("task_id")
            or kwargs.get("difficulty")
            or "easy"
        )
        self._difficulty = actual_difficulty
        tasks = TASK_SETS.get(actual_difficulty, TASK_SETS["easy"])
        self._current_task = random.choice(tasks)
        self._step_count = 0
        self._done = False
        self._last_reward = 0.01
        self._rewards_history = []
        self._stagnation_count = 0
        return self._make_observation(previous_attempt=None, feedback=None)

    def step(self, action: SQLAction) -> SQLObservation:
        if self._current_task is None:
            self.reset()

        self._step_count += 1
        reward_obj = grade(action, self._current_task)
        reward = reward_obj.value

        # Stagnation penalty: penalize repeating the same score
        if abs(reward - self._last_reward) < 0.01 and self._step_count > 1:
            self._stagnation_count += 1
            if self._stagnation_count >= 2:
                reward = max(0.01, reward - 0.1)
        else:
            self._stagnation_count = 0
        # Final Clamp  
        reward = max(0.01, min(0.98, reward))
          
        self._last_reward = reward

        self._rewards_history.append(reward)

        done = (reward >= 0.95) or (
            self._step_count >= self._current_task.max_steps
        )
        self._done = done
        feedback = generate_feedback(action, self._current_task, reward)

        return self._make_observation(
            previous_attempt=action.corrected_query,
            feedback=feedback,
        )

    @property
    def state(self) -> SQLState:
        if self._current_task is None:
            return SQLState(
                task_id="none",
                difficulty="none",
                step_count=0,
                max_steps=0,
                done=False,
                last_reward=0.01,
                rewards_history=[],
            )
        return SQLState(
            task_id=self._current_task.task_id,
            difficulty=self._difficulty,
            step_count=self._step_count,
            max_steps=self._current_task.max_steps,
            done=self._done,
            last_reward=self._last_reward,
            rewards_history=self._rewards_history,
        )

    # ── Internal helpers ──────────────────────────────────────────────────────

    def _make_observation(
        self,
        previous_attempt: str | None,
        feedback: str | None,
    ) -> SQLObservation:
        assert self._current_task is not None
        max_steps = self._current_task.max_steps
        steps_remaining = max(0, max_steps - self._step_count)
        return SQLObservation(
            task_id=self._current_task.task_id,
            broken_query=self._current_task.broken_query,
            schema_context=self._current_task.schema_context,
            # Only surface the hint on easy tasks
            error_hint=(
                self._current_task.error_hint
                if self._difficulty == "easy"
                else None
            ),
            step_number=self._step_count,
            steps_remaining=steps_remaining,
            previous_attempt=previous_attempt,
            feedback=feedback,
        )


app = create_app(
    SQLCorrectionEnvironment,
    SQLAction,
    SQLObservation,
    env_name="sql-correction-env",
)


@app.get("/tasks")
async def list_tasks():
    """List available task difficulties with metadata."""
    return {
        "tasks": [
            {
                "id": "easy",
                "difficulty": "easy",
                "description": "Fix a single SQL keyword typo. Error hint provided.",
                "max_steps": 5,
                "count": 15,
            },
            {
                "id": "medium",
                "difficulty": "medium",
                "description": (
                    "Fix multiple errors across keywords and clauses. No hint."
                ),
                "max_steps": 5,
                "count": 15,
            },
            {
                "id": "hard",
                "difficulty": "hard",
                "description": (
                    "Fix complex multi-join queries including column name errors. "
                    "Schema provided, no hint."
                ),
                "max_steps": 4,
                "count": 10,
            },
        ]
    }


def main():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)


if __name__ == "__main__":
    main()