openenv-testing/android_env-pr-162

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"""8HTTP server wrapper for Environment instances.9 10This module provides utilities to wrap any Environment subclass and expose it11over HTTP endpoints that HTTPEnvClient can consume.12"""13 14from __future__ import annotations15 16import os17from dataclasses import asdict18from typing import Any, Dict, Type19 20from .interfaces import Environment21from .types import Action, Observation22from fastapi import Body, FastAPI23 24class HTTPEnvServer:25    """26    HTTP server wrapper for Environment instances.27 28    This class wraps an Environment and exposes its reset(), step(), and state29    methods as HTTP endpoints compatible with HTTPEnvClient.30 31    The server expects:32    - Action deserialization: Converts JSON dict to Action subclass33    - Observation serialization: Converts Observation subclass to JSON dict34 35    Example:36        >>> from core.env_server import HTTPEnvServer37        >>> from envs.coding_env.server import CodeExecutionEnvironment38        >>>39        >>> env = CodeExecutionEnvironment()40        >>> server = HTTPEnvServer(env)41        >>>42        >>> # Register routes with FastAPI43        >>> from fastapi import FastAPI44        >>> app = FastAPI()45        >>> server.register_routes(app)46    """47 48    def __init__(49        self,50        env: Environment,51        action_cls: Type[Action],52        observation_cls: Type[Observation],53    ):54        """55        Initialize HTTP server wrapper.56 57        Args:58            env: The Environment instance to wrap59            action_cls: The Action subclass this environment expects60            observation_cls: The Observation subclass this environment returns61        """62        self.env = env63        self.action_cls = action_cls64        self.observation_cls = observation_cls65 66    def register_routes(self, app: Any) -> None:67        """68        Register HTTP routes on a FastAPI application.69 70        Args:71            app: FastAPI application instance72        """73 74        if not isinstance(app, FastAPI):75            raise TypeError("app must be a FastAPI instance")76 77        @app.post("/reset")78        async def reset(request: Dict[str, Any] = Body(default={})) -> Dict[str, Any]:79            """Reset endpoint - returns initial observation."""80            # TODO: Handle seed, episode_id from request if provided81            observation = self.env.reset()82            return self._serialize_observation(observation)83 84        @app.post("/step")85        async def step(request: Dict[str, Any]) -> Dict[str, Any]:86            """Step endpoint - executes action and returns observation."""87            action_data = request.get("action", {})88            # TODO: Handle timeout_s, request_id, episode_id from request if provided89 90            # Deserialize action91            action = self._deserialize_action(action_data)92 93            # Execute step94            observation = self.env.step(action)95 96            # Return serialized observation97            return self._serialize_observation(observation)98 99        @app.get("/state")100        async def get_state() -> Dict[str, Any]:101            """State endpoint - returns current environment state."""102            state = self.env.state103            return asdict(state)104 105        @app.get("/health")106        async def health() -> Dict[str, str]:107            """Health check endpoint."""108            return {"status": "healthy"}109 110 111    def _deserialize_action(self, action_data: Dict[str, Any]) -> Action:112        """113        Convert JSON dict to Action instance.114 115        Args:116            action_data: Dictionary containing action data117 118        Returns:119            Action instance120 121        Note:122            This is a simple implementation. Subclasses may need to override123            for more complex deserialization logic.124        """125        # Remove metadata if present (it will be set via kw_only field)126        metadata = action_data.pop("metadata", {})127        action = self.action_cls(**action_data)128        action.metadata = metadata129        return action130 131    def _serialize_observation(self, observation: Observation) -> Dict[str, Any]:132        """133        Convert Observation instance to JSON-compatible dict.134 135        Args:136            observation: Observation instance137 138        Returns:139            Dictionary compatible with HTTPEnvClient._parse_result()140 141        The format matches what HTTPEnvClient expects:142        {143            "observation": {...},  # Observation fields144            "reward": float | None,145            "done": bool,146        }147        """148        obs_dict = asdict(observation)149 150        # Extract reward and done (these are part of StepResult on client side)151        reward = obs_dict.pop("reward", None)152        done = obs_dict.pop("done", False)153        obs_dict.pop("metadata", None)  # Remove metadata from observation154 155        # Return in HTTPEnvClient expected format156        return {157            "observation": obs_dict,158            "reward": reward,159            "done": done,160        }161 162def create_app(163    env: Environment,164    action_cls: Type[Action],165    observation_cls: Type[Observation],166    env_name: Optional[str] = None,167) -> Any:168    """169    Create a FastAPI application with or without web interface.170    171    This function creates a FastAPI app with the web interface enabled by default,172    including README integration for better user experience.173    174    Args:175        env: The Environment instance to serve176        action_cls: The Action subclass this environment expects177        observation_cls: The Observation subclass this environment returns178        env_name: Optional environment name for README loading179        180    Returns:181        FastAPI application instance with or without web interface and README integration182    """183    # Check if web interface should be enabled184    # This can be controlled via environment variable or build argument185    enable_web = (186        os.getenv("ENABLE_WEB_INTERFACE", "false").lower() in ("true", "1", "yes")187    )188 189    if enable_web:190        # Import web interface only when needed191        from .web_interface import create_web_interface_app192        return create_web_interface_app(env, action_cls, observation_cls, env_name)193    else:194        # Use standard FastAPI app without web interface195        return create_fastapi_app(env, action_cls, observation_cls)196    197 198def create_fastapi_app(199    env: Environment,200    action_cls: Type[Action],201    observation_cls: Type[Observation],202) -> Any:203    """204    Create a FastAPI application with routes for the given environment.205 206    Args:207        env: The Environment instance to serve208        action_cls: The Action subclass this environment expects209        observation_cls: The Observation subclass this environment returns210 211    Returns:212        FastAPI application instance with routes registered213 214    Example:215        >>> from envs.coding_env.server import CodeExecutionEnvironment216        >>> from envs.coding_env.models import CodeAction, CodeObservation217        >>>218        >>> env = CodeExecutionEnvironment()219        >>> app = create_fastapi_app(env, CodeAction, CodeObservation)220        >>>221        >>> # Run with: uvicorn module:app --host 0.0.0.0 --port 8000222    """223    try:224        from fastapi import FastAPI225    except ImportError:226        raise ImportError(227            "FastAPI is required. Install with: pip install fastapi uvicorn"228        )229 230    app = FastAPI(title="Environment HTTP Server")231    server = HTTPEnvServer(env, action_cls, observation_cls)232    server.register_routes(app)233    return app234