client.py

"""Spreadsheet Environment Client.

Connects to a running Spreadsheet OpenEnv server over HTTP/WebSocket.
Agents interact via MCP tools (read_range, write_cell, submit_workbook, etc.).
"""

from __future__ import annotations

from typing import Any, Dict

from openenv.core.client_types import StepResult
from openenv.core.env_client import EnvClient
from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction, Tool

from .models import SpreadsheetAction, SpreadsheetObservation, SpreadsheetState


class SpreadsheetEnv(
    EnvClient[SpreadsheetAction, SpreadsheetObservation, SpreadsheetState]
):
    """Client for the Spreadsheet Environment.

    Example:
        >>> with SpreadsheetEnv(base_url="http://localhost:8000") as client:
        ...     result = client.reset()
        ...     result = client.step(
        ...         SpreadsheetAction(tool_name="list_scenarios", arguments_json="{}")
        ...     )
        ...     result = client.step(
        ...         SpreadsheetAction(
        ...             tool_name="read_range",
        ...             arguments_json='{"sheet":"Summary","range":"A1:D10"}'
        ...         )
        ...     )
    """

    def list_tools(self, use_cache: bool = True):
        if use_cache and hasattr(self, "_tools_cache") and self._tools_cache:
            return self._tools_cache
        import requests

        http_base = (
            self._ws_url
            .replace("ws://", "http://")
            .replace("wss://", "https://")
            .rstrip("/ws")
        )
        resp = requests.post(
            f"{http_base}/step",
            json={"action": {"type": "list_tools"}},
        )
        data = resp.json()
        raw_tools = data.get("observation", {}).get("tools", [])
        tools = [
            Tool(
                name=t["name"],
                description=t.get("description", ""),
                input_schema=t.get("input_schema", {}),
            )
            for t in raw_tools
        ]
        self._tools_cache = tools
        return tools

    def _step_payload(self, action: Any) -> Dict:
        if hasattr(action, "to_mcp_action"):
            action = action.to_mcp_action()
        if isinstance(action, ListToolsAction):
            return {"type": "list_tools"}
        if isinstance(action, CallToolAction):
            return {
                "type": "call_tool",
                "tool_name": action.tool_name,
                "arguments": action.arguments or {},
            }
        if hasattr(action, "model_dump"):
            return action.model_dump()
        return {"tool_name": getattr(action, "tool_name", ""), "arguments": {}}

    def _parse_result(self, payload: Dict) -> StepResult[SpreadsheetObservation]:
        obs_data = payload.get("observation", payload)
        observation = SpreadsheetObservation(
            tool_name=obs_data.get("tool_name", ""),
            result=obs_data.get("result"),
            error=obs_data.get("error"),
            done=payload.get("done", False),
            reward=payload.get("reward"),
            metadata=obs_data.get("metadata", {}),
        )
        return StepResult(
            observation=observation,
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> SpreadsheetState:
        return SpreadsheetState(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
        )