revert to working server

sql_env/server.py

@@ -1,114 +1,120 @@
 """
-FastAPI HTTP wrapper for SQLCorrectionEnv using openenv.core base classes.
+FastAPI HTTP wrapper for SQLCorrectionEnv.
+
+Exposes the OpenEnv-required endpoints: /reset, /step, /state + /tasks for validator.
 """
-from openenv.core.env_server import create_fastapi_app
-from openenv.core.env_server.interfaces import Environment
-from openenv.core.env_server.types import State
-from sql_env.models import SQLAction, SQLObservation
-from sql_env.tasks import ALL_TASKS, TASK_SETS
-from sql_env.grader import grade, generate_feedback
-import random
-
-
-class SQLCorrectionEnvironment(Environment):
-
-    def __init__(self):
-        super().__init__()
-        self._difficulty = "easy"
-        self._current_task = None
-        self._step_count = 0
-        self._done = False
-        self._last_reward = 0.0
-        self._previous_attempt = None
-        self._feedback = None
-        self._rewards_history = []
-
-    def reset(self, difficulty: str = "easy") -> SQLObservation:
-        self._difficulty = difficulty
-        tasks = TASK_SETS.get(difficulty, TASK_SETS["easy"])
-        self._current_task = random.choice(tasks)
-        self._step_count = 0
-        self._done = False
-        self._last_reward = 0.0
-        self._previous_attempt = None
-        self._feedback = None
-        self._rewards_history = []
-        return SQLObservation(
-            task_id=self._current_task.task_id,
-            broken_query=self._current_task.broken_query,
-            schema_context=self._current_task.schema_context,
-            error_hint=self._current_task.error_hint,
-            step_number=0,
-            previous_attempt=None,
-            feedback=None,
-        )
-
-    def step(self, action: SQLAction) -> SQLObservation:
-        self._step_count += 1
-        reward_obj = grade(action, self._current_task)
-        reward = reward_obj.value
-        self._last_reward = reward
-        self._previous_attempt = action.corrected_query
-        self._feedback = generate_feedback(action, self._current_task, reward_obj)
-        self._rewards_history.append(reward)
-        done = (reward >= 0.95) or (self._step_count >= self._current_task.max_steps)
-        self._done = done
-        return SQLObservation(
-            task_id=self._current_task.task_id,
-            broken_query=self._current_task.broken_query,
-            schema_context=self._current_task.schema_context,
-            error_hint=self._current_task.error_hint,
-            step_number=self._step_count,
-            previous_attempt=self._previous_attempt,
-            feedback=self._feedback,
-        )
-
-    @property
-    def state(self) -> dict:
-        if self._current_task is None:
-            return {"status": "not_initialized"}
-        return {
-            "task_id": self._current_task.task_id,
-            "difficulty": self._difficulty,
-            "step_count": self._step_count,
-            "done": self._done,
-            "last_reward": self._last_reward,
-            "rewards_history": self._rewards_history,
-        }
-
-    def get_tasks(self):
-        return {
-            "tasks": [
-                {
-                    "name": "easy",
-                    "difficulty": "easy",
-                    "description": "Fix a single syntax error. Error hint provided.",
-                    "max_steps": 5,
-                    "has_grader": True,
-                    "grader": "sql_env.grader.grade",
-                },
-                {
-                    "name": "medium",
-                    "difficulty": "medium",
-                    "description": "Fix multiple errors. No hint.",
-                    "max_steps": 5,
-                    "has_grader": True,
-                    "grader": "sql_env.grader.grade",
-                },
-                {
-                    "name": "hard",
-                    "difficulty": "hard",
-                    "description": "Fix complex multi-join queries. Schema provided.",
-                    "max_steps": 4,
-                    "has_grader": True,
-                    "grader": "sql_env.grader.grade",
-                },
-            ]
-        }
-
-
-env = SQLCorrectionEnvironment()
-app = create_fastapi_app(env, SQLAction, SQLObservation)
+
+from contextlib import asynccontextmanager
+from typing import Optional
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+
+from sql_env import SQLAction, SQLCorrectionEnv
+from sql_env.tasks import ALL_TASKS
+
+
+class ResetRequest(BaseModel):
+    difficulty: Optional[str] = "easy"
+    task_name: Optional[str] = None
+    task_index: Optional[int] = None
+
+
+class StepRequest(BaseModel):
+    corrected_query: str
+
+
+env: Optional[SQLCorrectionEnv] = None
+
+
+@asynccontextmanager
+async def lifespan(_: FastAPI):
+    global env
+    env = SQLCorrectionEnv(difficulty="easy")
+    yield
+    if env is not None:
+        await env.close()
+
+
+app = FastAPI(
+    title="SQL Correction RL Environment",
+    description="OpenEnv-compliant environment for SQL query correction tasks.",
+    version="1.0.0",
+    lifespan=lifespan,
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.post("/reset")
+async def reset(request: ResetRequest = ResetRequest()):
+    """Reset the environment and return the initial observation."""
+    global env
+    difficulty = request.task_name or request.difficulty or "easy"
+    if difficulty not in {"easy", "medium", "hard"}:
+        raise HTTPException(status_code=400, detail="difficulty must be easy, medium, or hard")
+
+    env = SQLCorrectionEnv(
+        difficulty=difficulty,
+        task_index=request.task_index,
+    )
+    obs = await env.reset()
+    return obs.model_dump()
+
+
+@app.post("/step")
+async def step(request: StepRequest):
+    """Take one step and return the new observation, reward, done flag, and info."""
+    global env
+    if env is None:
+        raise HTTPException(status_code=400, detail="Call /reset first.")
+
+    try:
+        action = SQLAction(corrected_query=request.corrected_query)
+        result = await env.step(action)
+        return result.model_dump()
+    except RuntimeError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+
+@app.post("/state")
+async def state():
+    """Return current environment state."""
+    global env
+    if env is None:
+        return {"status": "not_initialized"}
+    return await env.state()
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "service": "sql-correction-env"}
+
+
+@app.get("/tasks")
+async def list_tasks():
+    """Return graded tasks by difficulty (RL validator format)."""
+    graded_tasks = {
+        diff: [task.__dict__ for task in tasks if task.grader is not None]
+        for diff, tasks in ALL_TASKS.items()
+    }
+    return graded_tasks  # {"easy": [tasks], "medium": [tasks], "hard": [tasks]}
+
+
+@app.get("/")
+async def root():
+    return {
+        "name": "SQL Correction RL Environment",
+        "version": "1.0.0",
+        "endpoints": ["/reset", "/step", "/state", "/health", "/tasks"],
+        "tasks": ["easy", "medium", "hard"],
+    }
 
 
 def main():