| |
| |
| |
| |
| |
|
|
| import httpx |
| from typing import Optional |
| from .models import Observation, StepResult, Action |
|
|
|
|
| class StudyEnvClient: |
| """ |
| HTTP client for the Study Optimizer environment server. |
| |
| Usage: |
| client = StudyEnvClient("http://localhost:8000") |
| obs = client.reset() |
| result = client.step(0) # 0=study |
| """ |
|
|
| def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 30.0): |
| self.base_url = base_url.rstrip("/") |
| self.timeout = timeout |
| |
| |
| self._client = httpx.Client(timeout=self.timeout) |
|
|
| def health(self) -> dict: |
| """Quick check to see if the server is up.""" |
| resp = self._client.get(f"{self.base_url}/health") |
| resp.raise_for_status() |
| return resp.json() |
|
|
| def reset(self, task_id: Optional[str] = None) -> Observation: |
| """Reset environment and get first observation.""" |
| payload = {} |
| if task_id: |
| payload["task_id"] = task_id |
|
|
| resp = self._client.post(f"{self.base_url}/reset", json=payload) |
| resp.raise_for_status() |
| data = resp.json() |
|
|
| |
| if "observation" in data: |
| return Observation(**data["observation"]) |
| return Observation(**data) |
|
|
| def step(self, action: int) -> StepResult: |
| """Take an action and get the result.""" |
| payload = {"action": action} |
| resp = self._client.post(f"{self.base_url}/step", json=payload) |
| resp.raise_for_status() |
| data = resp.json() |
| return StepResult(**data) |
|
|
| def get_state(self) -> dict: |
| """Get the full current state of the environment.""" |
| resp = self._client.get(f"{self.base_url}/state") |
| resp.raise_for_status() |
| return resp.json() |
|
|
| def close(self): |
| """Clean up the HTTP client.""" |
| self._client.close() |
|
|
| def __enter__(self): |
| return self |
|
|
| def __exit__(self, *args): |
| self.close() |
|
|
