use openenv core base classes

sql_env/server.py

@@ -1,142 +1,114 @@
 """
-FastAPI HTTP wrapper for SQLCorrectionEnv.
-
-Exposes the OpenEnv-required endpoints: /reset, /step, /state + /tasks for validator.
+FastAPI HTTP wrapper for SQLCorrectionEnv using openenv.core base classes.
 """
-
-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 {
-        "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",
-            },
-        ]
-    }
-
-
-@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"],
-    }
+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)
 
 
 def main():