Spaces:
Sleeping
Sleeping
File size: 3,936 Bytes
1a1713a c4a66be 1a1713a c4a66be 6c6f994 c4a66be 6c6f994 c4a66be 6c6f994 c4a66be 6c6f994 c4a66be 1a1713a 5016a17 | 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 | """
FastAPI HTTP wrapper for SQLCorrectionEnv.
Exposes the OpenEnv-required endpoints: /reset, /step, /state + /tasks for validator.
"""
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():
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
if __name__ == "__main__":
main()
@app.post("/grader")
async def grader_endpoint(request: dict):
"""Grader endpoint called by validator to score a task directly."""
from sql_env.grader import grade
from sql_env.tasks import TASK_SETS
import random
task_name = request.get("task_name", "easy")
action_data = request.get("action", {})
corrected_query = action_data.get("corrected_query", "")
tasks = TASK_SETS.get(task_name, TASK_SETS["easy"])
task = random.choice(tasks)
action = SQLAction(corrected_query=corrected_query)
reward_obj = grade(action, task)
return {
"task_name": task_name,
"score": reward_obj.value,
"reason": reward_obj.reason,
"success": reward_obj.value >= 0.95,
}
|