File size: 2,348 Bytes
64b4cae | 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 | # client.py
# A simple Python client for calling the environment server
# You could do this with httpx directly but wrapping it is cleaner
#
# Use this in inference.py or anywhere else you need to talk to the server
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
# using httpx because requests doesn't support async
# TODO: add async versions of these methods later
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()
# took me a while to figure this out - the server wraps in 'observation' key
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()
|