models.py

"""Data models for the Spreadsheet Environment.

SpreadsheetAction has explicit Pydantic fields for MCP-style tool calls
(tool_name, arguments_json) compatible with the OpenEnv web interface.
"""

from __future__ import annotations

import json as _json
from typing import Any, Union

from pydantic import ConfigDict, Field, TypeAdapter, model_validator

from openenv.core.env_server.mcp_types import (
    CallToolAction,
    CallToolObservation,
    ListToolsAction,
    ListToolsObservation,
)
from openenv.core.env_server.types import Action, Observation, State

_mcp_action_adapter = TypeAdapter(Union[ListToolsAction, CallToolAction])

_AVAILABLE_TOOLS = (
    "list_tools, get_session_info, list_scenarios, load_scenario, "
    "list_sheets, read_range, write_cell, write_range, inspect_formula, "
    "list_named_targets, validate_partial, submit_workbook, "
    "get_edit_history, reset_scenario"
)


class SpreadsheetAction(Action):
    """Action with explicit fields for the web UI and MCP compatibility."""

    model_config = ConfigDict(
        extra="forbid",
        validate_assignment=True,
        arbitrary_types_allowed=True,
    )

    tool_name: str = Field(
        default="list_tools",
        description=f"MCP tool to invoke. Available: {_AVAILABLE_TOOLS}",
    )
    arguments_json: str = Field(
        default="{}",
        description=(
            'Tool arguments as a JSON string. Examples: '
            '"{}" for no args, '
            '\'{"scenario_id":"formula_repair_01"}\' for load_scenario, '
            '\'{"sheet":"Summary","range":"A1:D10"}\' for read_range, '
            '\'{"sheet":"Summary","cell":"C15","value":"=SUM(A1:A10)"}\' for write_cell'
        ),
    )

    @model_validator(mode="after")
    def _validate_json(self) -> "SpreadsheetAction":
        if self.arguments_json.strip():
            _json.loads(self.arguments_json)
        return self

    @classmethod
    def model_validate(cls, data: Any, **kwargs: Any) -> Action:
        if isinstance(data, dict) and data.get("type") in ("call_tool", "list_tools"):
            return _mcp_action_adapter.validate_python(data)
        return super().model_validate(data, **kwargs)

    def to_mcp_action(self) -> Action:
        if self.tool_name == "list_tools":
            return ListToolsAction()
        args = _json.loads(self.arguments_json) if self.arguments_json else {}
        return CallToolAction(tool_name=self.tool_name, arguments=args)


SpreadsheetObservation = CallToolObservation
SpreadsheetState = State

__all__ = [
    "SpreadsheetAction",
    "SpreadsheetObservation",
    "SpreadsheetState",
    "CallToolAction",
    "CallToolObservation",
    "ListToolsAction",
    "ListToolsObservation",
]