Spaces:
Sleeping
Sleeping
| """Typed HTTP client for the DB-ER environment.""" | |
| from __future__ import annotations | |
| from typing import Any, Dict, Optional | |
| from openenv.core import EnvClient | |
| from openenv.core.client_types import StepResult | |
| from db_er.models import DBERAction, DBERObservation, DBERState | |
| class DBERClient(EnvClient[DBERAction, DBERObservation, DBERState]): | |
| """Async client that connects to a running DB-ER server. | |
| Usage (async):: | |
| async with DBERClient(base_url="http://localhost:8000") as client: | |
| result = await client.reset() | |
| print(result.observation.incident_ticket) | |
| Usage (sync):: | |
| with DBERClient(base_url="http://localhost:8000").sync() as client: | |
| result = client.reset() | |
| obs = result.observation | |
| The ``reset()`` method accepts an optional ``task_id`` keyword argument | |
| (1=Easy, 2=Medium, 3=Hard). If omitted the server cycles tasks based on | |
| the provided ``seed``. | |
| """ | |
| def _step_payload(self, action: DBERAction) -> Dict[str, Any]: | |
| return action.model_dump(exclude_none=True) | |
| def _parse_result(self, payload: Dict[str, Any]) -> StepResult[DBERObservation]: | |
| obs_data = payload.get("observation", {}) | |
| reward = payload.get("reward") | |
| done = payload.get("done", False) | |
| obs = DBERObservation(**obs_data) | |
| return StepResult(observation=obs, reward=reward, done=done) | |
| def _parse_state(self, payload: Dict[str, Any]) -> DBERState: | |
| return DBERState(**payload) | |