naivaidhya/sql-data-quality-env

1"""2client.py — Python client for the SQL Data Quality Environment.3 4Provides a synchronous client (SQLDataQualityClient) that wraps the HTTP API,5following the OpenEnv EnvClient pattern.6 7Example:8    from client import SQLDataQualityClient, Action, ActionType9 10    with SQLDataQualityClient(base_url="http://localhost:7860") as client:11        obs = client.reset(task_id="easy")12        result = client.step(Action(action_type=ActionType.LIST_TABLES))13        state = client.get_state()14"""15 16from __future__ import annotations17 18from typing import Optional19 20import requests21 22from models import Action, Observation, State, StepResult23 24 25class SQLDataQualityClient:26    """27    Synchronous HTTP client for the SQL Data Quality Environment.28 29    Compatible with the OpenEnv EnvClient interface.30    """31 32    def __init__(self, base_url: str = "http://localhost:7860"):33        self.base_url = base_url.rstrip("/")34        self._session = requests.Session()35 36    def __enter__(self):37        return self38 39    def __exit__(self, *args):40        self._session.close()41 42    # ------------------------------------------------------------------43    # OpenEnv API44    # ------------------------------------------------------------------45 46    def reset(47        self,48        task_id: str = "easy",49        seed: Optional[int] = None,50        episode_id: Optional[str] = None,51    ) -> Observation:52        """Reset the environment and start a new episode."""53        payload = {"task_id": task_id}54        if seed is not None:55            payload["seed"] = seed56        if episode_id is not None:57            payload["episode_id"] = episode_id58        resp = self._session.post(f"{self.base_url}/reset", json=payload, timeout=30)59        resp.raise_for_status()60        return Observation(**resp.json())61 62    def step(self, action: Action) -> StepResult:63        """Execute one action and return the step result."""64        resp = self._session.post(65            f"{self.base_url}/step",66            json=action.model_dump(exclude_none=True),67            timeout=30,68        )69        resp.raise_for_status()70        data = resp.json()71        return StepResult(72            observation=Observation(**data["observation"]),73            reward=data["reward"],74            done=data["done"],75            info=data.get("info", {}),76        )77 78    def get_state(self) -> State:79        """Return the current episode state."""80        resp = self._session.get(f"{self.base_url}/state", timeout=10)81        resp.raise_for_status()82        return State(**resp.json())83 84    def health(self) -> dict:85        """Check server health."""86        resp = self._session.get(f"{self.base_url}/health", timeout=5)87        resp.raise_for_status()88        return resp.json()89 90    def list_tasks(self) -> dict:91        """Return all available tasks."""92        resp = self._session.get(f"{self.base_url}/tasks", timeout=10)93        resp.raise_for_status()94        return resp.json()95