openenv/tbench2

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"""TB2 Environment Client."""8 9from __future__ import annotations10 11from typing import Any12 13 14# Support both in-repo and standalone imports15try:16    # In-repo imports (when running from OpenEnv repository)17    from openenv.core.client_types import StepResult18    from openenv.core.env_client import EnvClient19 20    from .models import Tbench2Action, Tbench2Observation, Tbench2State21except ImportError:22    from models import Tbench2Action, Tbench2Observation, Tbench2State23 24    # Standalone imports (when environment is standalone with openenv from pip)25    from openenv.core.client_types import StepResult26    from openenv.core.env_client import EnvClient27 28 29class Tbench2Env(EnvClient[Tbench2Action, Tbench2Observation, Tbench2State]):30    """HTTP client for the TB2 environment."""31 32    def _step_payload(self, action: Tbench2Action) -> dict[str, Any]:33        return {34            "action_type": action.action_type,35            "command": action.command,36            "session_id": action.session_id,37            "block": action.block,38            "wait_seconds": action.wait_seconds,39            "file_path": action.file_path,40            "content": action.content,41        }42 43    def _parse_result(self, payload: dict[str, Any]) -> StepResult[Tbench2Observation]:44        obs_data = payload.get("observation", {})45        observation = Tbench2Observation(46            instruction=obs_data.get("instruction", ""),47            output=obs_data.get("output", ""),48            success=obs_data.get("success", True),49            error=obs_data.get("error", ""),50            task_id=obs_data.get("task_id", ""),51            task_path=obs_data.get("task_path", ""),52            session_id=obs_data.get("session_id"),53            action_type=obs_data.get("action_type", ""),54            info=obs_data.get("info", {}),55            reward=payload.get("reward"),56            done=payload.get("done", False),57            metadata=obs_data.get("metadata", {}),58        )59        return StepResult(60            observation=observation,61            reward=payload.get("reward"),62            done=payload.get("done", False),63        )64 65    def _parse_state(self, payload: dict[str, Any]) -> Tbench2State:66        return Tbench2State(67            episode_id=payload.get("episode_id"),68            step_count=payload.get("step_count", 0),69            task_id=payload.get("task_id", ""),70            task_path=payload.get("task_path", ""),71            terminal_ready=payload.get("terminal_ready", False),72            last_action_type=payload.get("last_action_type", ""),73            last_command=payload.get("last_command", ""),74            last_output=payload.get("last_output", ""),75        )76