albert-einstein-09/codedark

1"""2CodeDark Client3 4HTTP client for interacting with CodeDark environment server.5Follows OpenEnv EnvClient pattern.6"""7 8from typing import Any, Dict, Optional9import requests10 11 12class CodeDarkEnv:13    """Client for CodeDark environment.14 15    Example usage:16        env = CodeDarkEnv("http://localhost:8000")17        obs = env.reset()18        print(f"Task: {obs['question']}")19 20        obs = env.step("run_python", "<code>result = df.shape</code>")21        print(f"Result: {obs['stdout']}")22 23        obs = env.step("submit_answer", "<answer>11.26</answer>")24        print(f"Reward: {obs['reward']}")25    """26 27    def __init__(self, base_url: str = "http://localhost:8000", timeout: int = 30):28        """Initialize client.29 30        Args:31            base_url: Server URL32            timeout: Request timeout in seconds33        """34        self.base_url = base_url.rstrip("/")35        self.timeout = timeout36        self._session = requests.Session()37 38    def reset(39        self, task_id: Optional[str] = None, seed: Optional[int] = None40    ) -> Dict[str, Any]:41        """Reset environment for a new episode.42 43        Args:44            task_id: Specific task to load (optional)45            seed: Random seed for task selection (optional)46 47        Returns:48            Initial observation dict49        """50        payload = {}51        if task_id is not None:52            payload["task_id"] = task_id53        if seed is not None:54            payload["seed"] = seed55 56        response = self._session.post(57            f"{self.base_url}/reset",58            json=payload if payload else None,59            timeout=self.timeout,60        )61        response.raise_for_status()62        return response.json()63 64    def step(self, tool: str, args: str = "") -> Dict[str, Any]:65        """Execute an action.66 67        Args:68            tool: Tool name (run_python, read_notes, save_note, clarify, submit_answer)69            args: Tool-specific arguments70 71        Returns:72            Observation dict73        """74        response = self._session.post(75            f"{self.base_url}/step",76            json={"tool": tool, "args": args},77            timeout=self.timeout,78        )79        response.raise_for_status()80        return response.json()81 82    def state(self) -> Dict[str, Any]:83        """Get current environment state.84 85        Returns:86            State dict87        """88        response = self._session.get(89            f"{self.base_url}/state",90            timeout=self.timeout,91        )92        response.raise_for_status()93        return response.json()94 95    def health(self) -> Dict[str, Any]:96        """Check server health.97 98        Returns:99            Health status dict100        """101        response = self._session.get(102            f"{self.base_url}/health",103            timeout=self.timeout,104        )105        response.raise_for_status()106        return response.json()107 108    def metadata(self) -> Dict[str, Any]:109        """Get environment metadata.110 111        Returns:112            Metadata dict113        """114        response = self._session.get(115            f"{self.base_url}/metadata",116            timeout=self.timeout,117        )118        response.raise_for_status()119        return response.json()120 121    def schema(self) -> Dict[str, Any]:122        """Get environment type schemas.123 124        Returns:125            Schema dict for action, observation, state126        """127        response = self._session.get(128            f"{self.base_url}/schema",129            timeout=self.timeout,130        )131        response.raise_for_status()132        return response.json()133 134    # Convenience methods for common tools135 136    def run_python(self, code: str) -> Dict[str, Any]:137        """Execute Python code.138 139        Args:140            code: Python code to execute141 142        Returns:143            Observation dict144        """145        return self.step("run_python", f"<code>{code}</code>")146 147    def read_notes(self) -> Dict[str, Any]:148        """Read all saved notes.149 150        Returns:151            Observation dict152        """153        return self.step("read_notes", "")154 155    def save_note(self, content: str) -> Dict[str, Any]:156        """Save a note.157 158        Args:159            content: Note content160 161        Returns:162            Observation dict163        """164        return self.step("save_note", content)165 166    def clarify(self, question: str) -> Dict[str, Any]:167        """Ask a clarifying question.168 169        Args:170            question: Clarifying question171 172        Returns:173            Observation dict174        """175        return self.step("clarify", f"<question>{question}</question>")176 177    def submit_answer(self, answer: Any) -> Dict[str, Any]:178        """Submit final answer.179 180        Args:181            answer: Answer value182 183        Returns:184            Final observation with reward185        """186        return self.step("submit_answer", f"<answer>{answer}</answer>")187 188    def close(self):189        """Close the session."""190        self._session.close()191 192    def __enter__(self):193        return self194 195    def __exit__(self, exc_type, exc_val, exc_tb):196        self.close()197        return False198