Hariprita/nl2sql-openenv

1from fastapi import FastAPI2from fastapi.responses import JSONResponse3from pydantic import BaseModel4from typing import Optional5import uvicorn6 7from .environment import SQLAgentEnvironment8from .tasks import TASKS9 10app = FastAPI(title="SQL Agent Environment", version="1.0.0")11 12_env = SQLAgentEnvironment()13 14class StepRequest(BaseModel):15    sql_query: str16 17class ResetRequest(BaseModel):18    seed: Optional[int] = None19 20@app.get("/health")21def health():22    return {"status": "ok", "environment": "sql-agent-env", "version": "1.0.0"}23 24@app.post("/reset")25def reset(req: ResetRequest = None):26    obs = _env.reset()27    return JSONResponse(content=obs)28 29@app.post("/step")30def step(req: StepRequest):31    obs = _env.step(req.sql_query)32    return JSONResponse(content=obs)33 34@app.get("/state")35def state():36    return JSONResponse(content=_env.state)37 38@app.get("/tasks")39def list_tasks():40    return JSONResponse(content={"tasks": TASKS})41 42@app.get("/")43def root():44    return {"message": "SQL Agent Environment. POST /reset to start."}45 46def main():47    uvicorn.run(app, host="0.0.0.0", port=7860)48 49if __name__ == "__main__":50    main()51