Spaces:
Sleeping
Sleeping
File size: 4,199 Bytes
1a1713a c4a66be 1a1713a c4a66be 55bb0ce c4a66be 55bb0ce c4a66be 1a1713a 95707b2 | 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 | """
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 openenv.core.env_server import Environment
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 in openenv validator format."""
from sql_env.grader import grade
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",
"grader_fn": grade.__module__ + "." + grade.__qualname__,
},
{
"name": "medium",
"difficulty": "medium",
"description": "Fix multiple errors. No hint.",
"max_steps": 5,
"has_grader": True,
"grader": "sql_env.grader.grade",
"grader_fn": grade.__module__ + "." + grade.__qualname__,
},
{
"name": "hard",
"difficulty": "hard",
"description": "Fix complex multi-join queries. Schema provided.",
"max_steps": 4,
"has_grader": True,
"grader": "sql_env.grader.grade",
"grader_fn": grade.__module__ + "." + grade.__qualname__,
},
]
}
@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() |