models.py

"""
SupplyChainEnv — OpenEnv MCP-compliant typed models.

MCP tool-calling action space for supply chain disruption management.
Agent manages a global network of ports, factories, warehouses, and
shipping routes. Disruptions happen (storms, strikes, factory fires)
and the agent must reroute shipments to minimize cost and delay.
"""

from typing import Any, Dict, List, Literal, Optional
from pydantic import ConfigDict, Field
from openenv.core.env_server.types import Action, Observation, State


class MCPAction(Action):
    """MCP tool-calling action."""
    action_type: Literal["ListToolsAction", "ToolCallAction"] = Field(
        ..., description="Type of action"
    )
    tool_name: Optional[str] = Field(None, description="Tool to call")
    arguments: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Tool arguments")


class SupplyChainAction(MCPAction):
    pass


class MCPObservation(Observation):
    """MCP observation with tool results."""
    success: bool = Field(True)
    error_message: Optional[str] = Field(None)
    tools_list: Optional[List[Dict[str, Any]]] = Field(None)
    tool_result: Optional[Dict[str, Any]] = Field(None)
    metadata: Dict[str, Any] = Field(default_factory=dict)
    done: bool = Field(False)
    reward: Optional[float] = Field(None)


class SupplyChainObservation(MCPObservation):
    pass


class SupplyChainState(State):
    """Full state — the global supply chain network."""
    day: int = Field(default=0)
    total_days: int = Field(default=30)
    budget_usd: float = Field(default=10_000_000.0)
    spent_usd: float = Field(default=0.0)
    revenue_usd: float = Field(default=0.0)
    shipments_delivered: int = Field(default=0)
    shipments_lost: int = Field(default=0)
    active_disruptions: List[Dict[str, Any]] = Field(default_factory=list)
    network: Dict[str, Any] = Field(default_factory=dict)
    pending_shipments: List[Dict[str, Any]] = Field(default_factory=list)
    difficulty: str = Field(default="medium")