Spaces:
Sleeping
Sleeping
| """HTTP client for the DNS-Env OpenEnv environment.""" | |
| from __future__ import annotations | |
| import requests | |
| from typing import Any | |
| class DNSEnvClient: | |
| """Simple HTTP client for the DNS environment server. | |
| Wraps the REST endpoints exposed by the FastAPI server so callers | |
| can interact with the environment from plain Python without manually | |
| assembling HTTP requests. | |
| Parameters | |
| ---------- | |
| base_url: | |
| Root URL of the environment server (default ``http://localhost:7860``). | |
| session_id: | |
| Identifies the caller's session. The server maintains independent | |
| environment state per session. | |
| timeout: | |
| HTTP request timeout in seconds. | |
| """ | |
| def __init__( | |
| self, | |
| base_url: str = "http://localhost:7860", | |
| session_id: str = "default", | |
| timeout: float = 30.0, | |
| ): | |
| self.base_url = base_url.rstrip("/") | |
| self.session_id = session_id | |
| self.timeout = timeout | |
| # ------------------------------------------------------------------ | |
| # Endpoints | |
| # ------------------------------------------------------------------ | |
| def health(self) -> dict: | |
| """Liveness check. Returns ``{"status": "ok"}`` when the server is up.""" | |
| resp = requests.get(f"{self.base_url}/health", timeout=self.timeout) | |
| resp.raise_for_status() | |
| return resp.json() | |
| def reset( | |
| self, | |
| task_id: str | None = None, | |
| seed: int | None = None, | |
| episode_id: str | None = None, | |
| ) -> dict: | |
| """Reset the environment and start a new episode. | |
| Parameters | |
| ---------- | |
| task_id: | |
| One of ``fix_single_record``, ``configure_mail``, | |
| ``debug_delegation``. When *None* the server cycles tasks. | |
| seed: | |
| Optional RNG seed for reproducibility. | |
| episode_id: | |
| Optional caller-supplied episode identifier. | |
| Returns | |
| ------- | |
| dict | |
| Observation JSON with keys: ``output``, ``task_description``, | |
| ``zone_names``, ``available_commands``, ``done``, ``reward``, | |
| ``metadata``. | |
| """ | |
| body: dict[str, Any] = {"session_id": self.session_id} | |
| options: dict[str, Any] = {} | |
| if task_id: | |
| options["task_id"] = task_id | |
| if options: | |
| body["options"] = options | |
| if seed is not None: | |
| body["seed"] = seed | |
| if episode_id is not None: | |
| body["episode_id"] = episode_id | |
| resp = requests.post( | |
| f"{self.base_url}/reset", json=body, timeout=self.timeout | |
| ) | |
| resp.raise_for_status() | |
| return resp.json() | |
| def step(self, command: str, **args: Any) -> dict: | |
| """Execute one action in the environment. | |
| Parameters | |
| ---------- | |
| command: | |
| One of the available commands (``view_zone``, ``add_record``, | |
| ``edit_record``, ``delete_record``, ``check_zone``, ``dig``, | |
| ``submit``). | |
| **args: | |
| Keyword arguments forwarded as the action's ``args`` dict. | |
| Returns | |
| ------- | |
| dict | |
| Observation JSON. | |
| """ | |
| body = { | |
| "session_id": self.session_id, | |
| "action": {"command": command, "args": args}, | |
| } | |
| resp = requests.post( | |
| f"{self.base_url}/step", json=body, timeout=self.timeout | |
| ) | |
| resp.raise_for_status() | |
| return resp.json() | |
| def state(self) -> dict: | |
| """Return the current episode state (step count, task id, etc.).""" | |
| resp = requests.get( | |
| f"{self.base_url}/state", | |
| params={"session_id": self.session_id}, | |
| timeout=self.timeout, | |
| ) | |
| resp.raise_for_status() | |
| return resp.json() | |
| def tasks(self) -> list[str]: | |
| """List available task identifiers.""" | |
| resp = requests.get(f"{self.base_url}/tasks", timeout=self.timeout) | |
| resp.raise_for_status() | |
| return resp.json()["tasks"] | |