Spaces:
Sleeping
Sleeping
| """ | |
| SQL Debug Environment — OpenEnv-compliant async client. | |
| Implements the standard EnvClient interface: | |
| async with SQLDebugEnv(base_url="...") as client: | |
| result = await client.reset() | |
| result = await client.step(SQLFixAction(query="SELECT ...")) | |
| Sync wrapper is also available: | |
| with SQLDebugEnv(base_url="...").sync() as client: | |
| result = client.reset() | |
| result = client.step(SQLFixAction(query="SELECT ...")) | |
| """ | |
| import asyncio | |
| from dataclasses import dataclass | |
| from typing import Optional | |
| import httpx | |
| from models import Action, Observation, State | |
| class StepResult: | |
| """Returned by step() and reset().""" | |
| observation: Observation | |
| reward: float | |
| done: bool | |
| success: bool | |
| state: Optional[State] = None | |
| class SQLDebugEnv: | |
| """ | |
| Async HTTP client for the SQL Debug Environment. | |
| Compatible with the OpenEnv EnvClient interface. | |
| Usage (async): | |
| async with SQLDebugEnv(base_url="http://localhost:8000") as env: | |
| result = await env.reset(task_group="easy_syntax_fix") | |
| result = await env.step(SQLFixAction(query="SELECT 1")) | |
| Usage (sync): | |
| with SQLDebugEnv(base_url="http://localhost:8000").sync() as env: | |
| result = env.reset() | |
| result = env.step(SQLFixAction(query="SELECT 1")) | |
| """ | |
| def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 30.0): | |
| self.base_url = base_url.rstrip("/") | |
| self.timeout = timeout | |
| self._session_id: Optional[str] = None | |
| self._client: Optional[httpx.AsyncClient] = None | |
| #Context Manager | |
| async def __aenter__(self): | |
| self._client = httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout) | |
| return self | |
| async def __aexit__(self, *args): | |
| if self._client: | |
| await self._client.aclose() | |
| self._client = None | |
| def sync(self) -> "SyncSQLDebugEnv": | |
| """Return a synchronous wrapper around this async client.""" | |
| return SyncSQLDebugEnv(self) | |
| #Core OpenEnv API | |
| async def reset( | |
| self, | |
| task_group: Optional[str] = None, | |
| task_id: Optional[str] = None, | |
| seed: Optional[int] = None, | |
| ) -> StepResult: | |
| """Initialize or restart an episode.""" | |
| payload = { | |
| "session_id": self._session_id, | |
| "task_group": task_group, | |
| "task_id": task_id, | |
| "seed": seed, | |
| } | |
| resp = await self._post("/reset", payload) | |
| self._session_id = resp["session_id"] | |
| obs = Observation(**resp["observation"]) | |
| state = State(**resp["state"]) if "state" in resp else None | |
| return StepResult(observation=obs, reward=0.0, done=False, success=False, state=state) | |
| async def step(self, action: Action) -> StepResult: | |
| """Submit a SQL fix attempt.""" | |
| if not self._session_id: | |
| raise RuntimeError("No active session. Call reset() first.") | |
| payload = { | |
| "session_id": self._session_id, | |
| "query": action.query, | |
| "explanation": action.explanation or "", | |
| } | |
| resp = await self._post("/step", payload) | |
| obs = Observation(**resp["observation"]) | |
| state = State(**resp["state"]) if "state" in resp else None | |
| return StepResult( | |
| observation=obs, | |
| reward=resp["reward"], | |
| done=resp["done"], | |
| success=resp["success"], | |
| state=state, | |
| ) | |
| async def state(self) -> State: | |
| """Return current episode state without acting.""" | |
| if not self._session_id: | |
| raise RuntimeError("No active session. Call reset() first.") | |
| client = self._client or httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout) | |
| try: | |
| r = await client.get("/state", params={"session_id": self._session_id}) | |
| r.raise_for_status() | |
| return State(**r.json()) | |
| finally: | |
| if not self._client: | |
| await client.aclose() | |
| async def _post(self, endpoint: str, payload: dict) -> dict: | |
| client = self._client | |
| own_client = False | |
| if client is None: | |
| client = httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout) | |
| own_client = True | |
| try: | |
| resp = await client.post(endpoint, json=payload) | |
| resp.raise_for_status() | |
| return resp.json() | |
| finally: | |
| if own_client: | |
| await client.aclose() | |
| def from_hub(cls, repo_id: str, port: int = 8000, **kwargs) -> "SQLDebugEnv": | |
| """ | |
| Connect to an environment hosted on Hugging Face Spaces. | |
| Example: | |
| env = SQLDebugEnv.from_hub("openenv/sql-debug-env") | |
| """ | |
| space_name = repo_id.replace("/", "-").lower() | |
| base_url = f"https://{space_name}.hf.space" | |
| return cls(base_url=base_url, **kwargs) | |
| class SyncSQLDebugEnv: | |
| """Synchronous wrapper around SQLDebugEnv.""" | |
| def __init__(self, async_env: SQLDebugEnv): | |
| self._env = async_env | |
| self._loop: Optional[asyncio.AbstractEventLoop] = None | |
| def __enter__(self): | |
| self._loop = asyncio.new_event_loop() | |
| self._loop.run_until_complete(self._env.__aenter__()) | |
| return self | |
| def __exit__(self, *args): | |
| self._loop.run_until_complete(self._env.__aexit__(*args)) | |
| self._loop.close() | |
| def reset(self, **kwargs) -> StepResult: | |
| return self._loop.run_until_complete(self._env.reset(**kwargs)) | |
| def step(self, action: Action) -> StepResult: | |
| return self._loop.run_until_complete(self._env.step(action)) | |
| def state(self) -> State: | |
| return self._loop.run_until_complete(self._env.state()) | |
| SQLFixAction = Action |