Spaces:
Sleeping
Sleeping
File size: 5,808 Bytes
be62603 | 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 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """
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
@dataclass
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()
@classmethod
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 |