huzzle-labs/visual_memory

1"""2Visual Memory Environment HTTP Client.3 4Connects to a running Visual Memory OpenEnv server over HTTP/WebSocket.5Agents interact via MCP tools exposed through step(CallToolAction(...)).6"""7 8from __future__ import annotations9 10from typing import Any, Dict11 12from openenv.core.client_types import StepResult13from openenv.core.env_client import EnvClient14from openenv.core.env_server.mcp_types import (15    CallToolAction,16    ListToolsAction,17    Tool,18)19 20from .models import (21    VisualMemoryAction,22    VisualMemoryObservation,23    VisualMemoryState,24)25 26 27class VisualMemoryEnv(EnvClient[VisualMemoryAction, VisualMemoryObservation, VisualMemoryState]):28    """HTTP client for the Visual Memory Environment.29 30    Example:31        >>> async with VisualMemoryEnv(base_url="http://localhost:8000") as client:32        ...     result = await client.reset()33        ...     result = await client.step(34        ...         CallToolAction(tool_name="load_scenario", arguments={"scenario_id": "hidden_grid_01"})35        ...     )36    """37 38    def list_tools(self, use_cache: bool = True):39        if use_cache and hasattr(self, "_tools_cache") and self._tools_cache:40            return self._tools_cache41        import requests42 43        http_base = (44            self._ws_url45            .replace("ws://", "http://")46            .replace("wss://", "https://")47            .rstrip("/ws")48        )49        resp = requests.post(50            f"{http_base}/step",51            json={"action": {"type": "list_tools"}},52        )53        data = resp.json()54        raw_tools = data.get("observation", {}).get("tools", [])55        tools = [56            Tool(57                name=t["name"],58                description=t.get("description", ""),59                input_schema=t.get("input_schema", {}),60            )61            for t in raw_tools62        ]63        self._tools_cache = tools64        return tools65 66    def _step_payload(self, action: Any) -> Dict:67        if isinstance(action, ListToolsAction):68            return {"type": "list_tools"}69        if isinstance(action, CallToolAction):70            return {71                "type": "call_tool",72                "tool_name": action.tool_name,73                "arguments": action.arguments or {},74            }75        if hasattr(action, "model_dump"):76            return action.model_dump()77        return {"tool_name": getattr(action, "tool_name", ""), "arguments": {}}78 79    def _parse_result(self, payload: Dict) -> StepResult[VisualMemoryObservation]:80        obs_data = payload.get("observation", payload)81        observation = VisualMemoryObservation(82            tool_name=obs_data.get("tool_name", ""),83            result=obs_data.get("result"),84            error=obs_data.get("error"),85            done=payload.get("done", False),86            reward=payload.get("reward"),87            metadata=obs_data.get("metadata", {}),88        )89        return StepResult(90            observation=observation,91            reward=payload.get("reward"),92            done=payload.get("done", False),93        )94 95    def _parse_state(self, payload: Dict[str, Any]) -> VisualMemoryState:96        return VisualMemoryState(97            episode_id=payload.get("episode_id"),98            step_count=payload.get("step_count", 0),99        )100