Team Ai
Apppublic

openenv/coding_env

sourceHugging Faceupdated 3mo agoView on Hugging Face
21likes
http_env_client.py186 linesDownload Raw Back to core
1"""2core/runner_env.py3Minimal HTTP-based environment client.4- Talks to a single env worker exposing: POST /reset, POST /step5 6Future hooks (commented below) for:7- episode_id, seed on reset8- request_id on step9- custom headers (auth/trace)10"""11 12from __future__ import annotations13 14from abc import ABC, abstractmethod15from typing import Any, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar16 17import requests18 19from .client_types import StepResult20from .containers.runtime import LocalDockerProvider21 22if TYPE_CHECKING:23    from .containers.runtime import ContainerProvider24 25ActT = TypeVar("ActT")26ObsT = TypeVar("ObsT")27EnvClientT = TypeVar("EnvClientT", bound="HTTPEnvClient")28 29 30class HTTPEnvClient(ABC, Generic[ActT, ObsT]):31    def __init__(32        self,33        base_url: str,34        request_timeout_s: float = 15.0,35        default_headers: Optional[Dict[str, str]] = None,36        provider: Optional["ContainerProvider"] = None,37    ):38        self._base = base_url.rstrip("/")39        self._timeout = float(request_timeout_s)40        self._http = requests.Session()41        self._headers = default_headers or {}42        self._provider = provider43 44    @classmethod45    def from_docker_image(46        cls: Type[EnvClientT],47        image: str,48        provider: Optional["ContainerProvider"] = None,49        **kwargs: Any,50    ) -> EnvClientT:51        """52        Create an environment client by spinning up a Docker container locally.53 54        This is a development utility that:55        1. Starts a Docker container from the specified image56        2. Waits for the server to be ready57        3. Creates and returns a client instance connected to the container58 59        Note: The container lifecycle management is left to the user or higher-level60        orchestration. The container will keep running until manually stopped.61 62        Args:63            image: Docker image name to run (e.g., "echo-env:latest")64            provider: Container provider to use (defaults to LocalDockerProvider)65            **kwargs: Additional arguments to pass to provider.start_container()66                     (e.g., env_vars, port)67 68        Returns:69            An instance of the client class connected to the running container70 71        Example:72            >>> from envs.coding_env.client import CodingEnv73            >>> from envs.coding_env.models import CodeAction74            >>>75            >>> # Create environment from image76            >>> env = CodingEnv.from_docker_image("coding-env:latest")77            >>>78            >>> # Create environment with custom env vars79            >>> env = CodingEnv.from_docker_image(80            ...     "coding-env:latest",81            ...     env_vars={"MY_VAR": "value"}82            ... )83            >>>84            >>> # Use the environment85            >>> result = env.reset()86            >>> print(result.observation)87            >>>88            >>> step_result = env.step(CodeAction(code="print('hello')"))89            >>> print(step_result.observation.stdout)90            >>>91            >>> # Cleanup (optional)92            >>> env.close()93        """94 95        # Use default provider if none provided96        if provider is None:97            provider = LocalDockerProvider()98 99        # 1. Start container with optional kwargs (e.g., env_vars, port)100        base_url = provider.start_container(image, **kwargs)101 102        # 2. Wait for server to be ready103        provider.wait_for_ready(base_url)104 105        # 3. Create and return client instance with provider reference106        return cls(base_url=base_url, provider=provider)107 108    @abstractmethod109    def _step_payload(self, action: ActT) -> dict:110        """Convert an Action object to the JSON body expected by the env server."""111        raise NotImplementedError112 113    @abstractmethod114    def _parse_result(self, payload: dict) -> StepResult[ObsT]:115        """Convert a JSON response from the env server to StepResult[ObsT]."""116        raise NotImplementedError117 118    @abstractmethod119    def _parse_state(self, payload: dict) -> Any:120        """Convert a JSON response from the state endpoint to a State object."""121        raise NotImplementedError122 123    # ---------- Environment Server Interface Methods ----------124    def reset(self) -> StepResult[ObsT]:125        body: Dict[str, Any] = {}126        # TODO: later:127        # body["seed"] = seed128        # body["episode_id"] = episode_id129        r = self._http.post(130            f"{self._base}/reset",131            json=body,132            headers=self._headers,133            timeout=self._timeout,134        )135        r.raise_for_status()136        return self._parse_result(r.json())137 138    def step(self, action: ActT) -> StepResult[ObsT]:139        body: Dict[str, Any] = {140            "action": self._step_payload(action),141            "timeout_s": int(self._timeout),142        }143        # TODO: later:144        # body["request_id"] = str(uuid.uuid4())145        # body["episode_id"] = current_episode_id146        r = self._http.post(147            f"{self._base}/step",148            json=body,149            headers=self._headers,150            timeout=self._timeout,151        )152        r.raise_for_status()153        return self._parse_result(r.json())154 155    def state(self) -> Any:156        """157        Get the current environment state from the server.158 159        Returns:160            State object with environment state information (e.g., episode_id, step_count)161 162        Example:163            >>> client = EchoEnv.from_docker_image("echo-env:latest")164            >>> result = client.reset()165            >>> state = client.state()166            >>> print(state.episode_id)167            >>> print(state.step_count)168        """169        r = self._http.get(170            f"{self._base}/state",171            headers=self._headers,172            timeout=self._timeout,173        )174        r.raise_for_status()175        return self._parse_state(r.json())176 177    def close(self) -> None:178        """179        Close the environment and clean up resources.180 181        If this client was created via from_docker_image(), this will stop182        and remove the associated container.183        """184        if self._provider is not None:185            self._provider.stop_container()186