Hariprita/nl2sql-openenv

1"""2client.py — SQLAgentEnv client class.3 4Mirrors the BrowserGymEnv pattern from the official OpenEnv sample inference script:5  env = SQLAgentEnv.from_docker_image(image="sql-agent-env:latest", env_vars={...})6  result = env.reset()                        # StepResult7  result = env.step(SQLAction(sql_query=...)) # StepResult8  env.close()9"""10 11import subprocess12import time13from dataclasses import dataclass, fields14from typing import Optional15 16import requests17 18 19# ---------------------------------------------------------------------------20# Data types21# ---------------------------------------------------------------------------22 23@dataclass24class SQLObservation:25    schema: str26    question: str27    result: str28    reward: float29    done: bool30    feedback: str31    task_id: str32    task_difficulty: str33    attempt: int34    max_attempts: int35    hint: str = ""36 37    @property38    def goal(self) -> str:39        """Alias so generic inference scripts can use observation.goal"""40        return self.question41 42    @property43    def last_action_error(self) -> bool:44        return bool(self.result and self.result.startswith("ERROR:"))45 46    @property47    def url(self) -> str:48        """Stub for compatibility with generic inference scripts"""49        return f"task://{self.task_id}/attempt/{self.attempt}"50 51 52@dataclass53class StepResult:54    observation: SQLObservation55    reward: float56    done: bool57 58 59@dataclass60class SQLAction:61    sql_query: str62 63 64# ---------------------------------------------------------------------------65# Environment client66# ---------------------------------------------------------------------------67 68class SQLAgentEnv:69    """70    Client for the SQL Agent Environment FastAPI server.71 72    Usage (local server already running):73        env = SQLAgentEnv(base_url="http://localhost:7860")74 75    Usage (spin up Docker image):76        env = SQLAgentEnv.from_docker_image("sql-agent-env:latest")77 78    Both return the same object; use reset() / step() / close() uniformly.79    """80 81    def __init__(self, base_url: str = "http://localhost:7860"):82        self.base_url = base_url.rstrip("/")83        self._container_id: Optional[str] = None84 85    # ------------------------------------------------------------------86    # Factory87    # ------------------------------------------------------------------88 89    @classmethod90    def from_docker_image(91        cls,92        image: str = "sql-agent-env:latest",93        port: int = 7860,94        env_vars: Optional[dict] = None,95    ) -> "SQLAgentEnv":96        """97        Launch the environment inside Docker, wait for it to be healthy,98        and return a connected client.99        """100        env_args: list = []101        for k, v in (env_vars or {}).items():102            env_args += ["-e", f"{k}={v}"]103 104        cmd = ["docker", "run", "-d", "--rm", "-p", f"{port}:{port}"] + env_args + [image]105        proc = subprocess.run(cmd, capture_output=True, text=True, check=True)106        container_id = proc.stdout.strip()107 108        client = cls(base_url=f"http://localhost:{port}")109        client._container_id = container_id110 111        # Poll until /health returns 200 (up to 60 s)112        for _ in range(60):113            try:114                r = requests.get(f"{client.base_url}/health", timeout=2)115                if r.status_code == 200:116                    return client117            except requests.exceptions.ConnectionError:118                pass119            time.sleep(1)120 121        raise TimeoutError(122            f"Environment did not become healthy within 60 s (container: {container_id})"123        )124 125    # ------------------------------------------------------------------126    # Core API127    # ------------------------------------------------------------------128 129    def reset(self) -> StepResult:130        r = requests.post(f"{self.base_url}/reset", json={}, timeout=10)131        r.raise_for_status()132        return self._parse_step_result(r.json())133 134    def step(self, action: SQLAction) -> StepResult:135        r = requests.post(136            f"{self.base_url}/step",137            json={"sql_query": action.sql_query},138            timeout=10,139        )140        r.raise_for_status()141        return self._parse_step_result(r.json())142 143    def state(self) -> dict:144        r = requests.get(f"{self.base_url}/state", timeout=10)145        r.raise_for_status()146        return r.json()147 148    # ------------------------------------------------------------------149    # Lifecycle150    # ------------------------------------------------------------------151 152    def close(self):153        if self._container_id:154            subprocess.run(155                ["docker", "stop", self._container_id],156                capture_output=True,157            )158            self._container_id = None159 160    def __enter__(self):161        return self162 163    def __exit__(self, *args):164        self.close()165 166    # ------------------------------------------------------------------167    # Internal helpers168    # ------------------------------------------------------------------169 170    @staticmethod171    def _parse_step_result(data: dict) -> StepResult:172        obs_field_names = {f.name for f in fields(SQLObservation)}173        obs = SQLObservation(**{k: data[k] for k in obs_field_names if k in data})174        return StepResult(175            observation=obs,176            reward=float(data.get("reward", 0.0)),177            done=bool(data.get("done", False)),178        )179