grid_env/tools.py

"""
Tool generation helpers for the warehouse fulfillment environment.
"""

from __future__ import annotations

from typing import Any, Dict, List

from .env import WarehouseFulfillmentEnv
from .models import WarehouseAction


_TOOL_DESCRIPTIONS: Dict[str, str] = {
    "turn_left": "Rotate the robot 90 degrees left in place.",
    "turn_right": "Rotate the robot 90 degrees right in place.",
    "move_forward": "Move one cell forward if not blocked.",
    "scan_bin": "Scan the bin in front to validate its contents.",
    "pick_item": "Pick an item from the bin in front.",
    "pack_item": "Pack the carried item at the packing station.",
    "recharge": "Recharge the battery at the charging dock.",
    "rest": "Rest at the rest area to restore stamina.",
    "wait": "Stay in place and consume time.",
}


def build_tools() -> List[Dict[str, Any]]:
    """Return OpenAI-compatible tool definitions for warehouse actions."""
    tools = []
    for command in WarehouseFulfillmentEnv.action_space:
        tools.append(_tool_for_command(command))
    return tools


def _tool_for_command(command: str) -> Dict[str, Any]:
    return {
        "type": "function",
        "function": {
            "name": command,
            "description": _TOOL_DESCRIPTIONS.get(command, "Warehouse action."),
            "parameters": {
                "type": "object",
                "properties": {},
                "additionalProperties": False,
            },
        },
    }


def tool_call_to_action(tool_call: Dict[str, Any]) -> WarehouseAction:
    """Convert a tool call payload into a WarehouseAction."""
    name = tool_call.get("name")
    if name is None and isinstance(tool_call.get("function"), dict):
        name = tool_call["function"].get("name")
    if name not in WarehouseFulfillmentEnv.action_space:
        raise ValueError(f"Unknown tool/action: {name}")
    return WarehouseAction(command=name)


def is_valid_tool_name(name: str) -> bool:
    return name in WarehouseFulfillmentEnv.action_space