Ar-Srivas/BitWise_CSS_env

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""Css Env Environment Client."""8 9from typing import Dict10 11from openenv.core import EnvClient12from openenv.core.client_types import StepResult13from openenv.core.env_server.types import State14 15try:16    from .models import CssAction, CssObservation17except ImportError:18    from models import CssAction, CssObservation19 20 21class CssEnv(22    EnvClient[CssAction, CssObservation, State]23):24    """25    Client for the Css Env Environment.26 27    This client maintains a persistent WebSocket connection to the environment server,28    enabling efficient multi-step interactions with lower latency.29    Each client instance has its own dedicated environment session on the server.30 31    Example:32        >>> # Connect to a running server33        >>> with CssEnv(base_url="http://localhost:8000") as client:34        ...     result = client.reset(task={"html": "<div class='card'></div>", "css": ".card{color:#1a6fe0;}", "tokens": {}, "config": {}}, seed=7)35        ...     print(result.observation.css)36        ...37        ...     result = client.step(CssAction(action_type="replace_color", target="#1a6fe0", value="#333333"))38        ...     print(result.observation.css)39 40    Example with Docker:41        >>> # Automatically start container and connect42        >>> client = CssEnv.from_docker_image("css_env-env:latest")43        >>> try:44        ...     result = client.reset(task={"html": "<div class='card'></div>", "css": ".card{color:#1a6fe0;}", "tokens": {}, "config": {}}, seed=7)45        ...     result = client.step(CssAction(action_type="remove_rule", target=".unused", value=None))46        ... finally:47        ...     client.close()48    """49 50    def _step_payload(self, action: CssAction) -> Dict:51        """52        Convert CssAction to JSON payload for step message.53 54        Args:55            action: CssAction instance56 57        Returns:58            Dictionary representation suitable for JSON encoding59        """60        return {61            "action_type": action.action_type,62            "target": action.target,63            "value": action.value,64        }65 66    def _parse_result(self, payload: Dict) -> StepResult[CssObservation]:67        """68        Parse server response into StepResult[CssObservation].69 70        Args:71            payload: JSON response data from server72 73        Returns:74            StepResult with CssObservation75        """76        obs_data = payload.get("observation", {})77        observation = CssObservation(78            html=obs_data.get("html", ""),79            css=obs_data.get("css", ""),80            tokens=obs_data.get("tokens", {}),81            violations=obs_data.get("violations"),82            scores=obs_data.get("scores"),83            score=obs_data.get("score"),84            success=obs_data.get("success"),85            changed=obs_data.get("changed"),86            no_op_action=obs_data.get("no_op_action"),87            repeated_action=obs_data.get("repeated_action"),88            terminated_by_max_steps=obs_data.get("terminated_by_max_steps"),89        )90 91        return StepResult(92            observation=observation,93            reward=payload.get("reward"),94            done=payload.get("done", False),95        )96 97    def _parse_state(self, payload: Dict) -> State:98        """99        Parse server response into State object.100 101        Args:102            payload: JSON response from state request103 104        Returns:105            State object with episode_id and step_count106        """107        return State(108            episode_id=payload.get("episode_id"),109            step_count=payload.get("step_count", 0),110        )111