ChilleD/agent_world_model_env

1"""2AWM-specific Pydantic models for action and observation types.3"""4 5from typing import Annotated, Any6 7from openenv.core.env_server.mcp_types import (8    CallToolAction,9    ListToolsAction,10    ListToolsObservation,11)12from openenv.core.env_server.types import Action, Observation13from pydantic import ConfigDict, Field, field_validator, TypeAdapter14 15 16_AWMActionUnion = Annotated[17    ListToolsAction | CallToolAction,18    Field(discriminator="type"),19]20_awm_action_adapter = TypeAdapter(_AWMActionUnion)21 22 23class AWMAction(Action):24    """Discriminated union action type for AWM.25 26    model_validate() returns the concrete ListToolsAction or CallToolAction27    (not an AWMAction instance), which is what AWMEnvironment.step() expects.28    """29 30    @classmethod31    def model_validate(cls, obj: Any, **kwargs: Any) -> Action:  # type: ignore[override]32        return _awm_action_adapter.validate_python(obj)33 34    @classmethod35    def model_json_schema(cls, **kwargs: Any) -> dict[str, Any]:  # type: ignore[override]36        return _awm_action_adapter.json_schema(**kwargs)37 38 39class AWMObservation(Observation):40    """41    Observation with AWM-specific fields promoted to top level.42    model_dump() excludes None-valued fields by default so that keys like43    ``tool_name=None`` do not appear in the wire payload.44    This is because the generic MCPToolClient._parse_result() routes observations based on key presence (e.g. ``"tool_name" in obs_data``). We may need to modify the MCPToolClient in the future. Currently, I try to avoid modifying any openenv code.45    """46 47    model_config = ConfigDict(extra="forbid")48 49    reward_type: str | None = Field(50        default=None,51        description="Reward classification label for this step/episode outcome",52    )53    scenario: str | None = Field(default=None, description="Current scenario name")54    task: str | None = Field(default=None, description="Current task description")55    task_idx: int | None = Field(default=None, description="Current task index")56    has_verifier: dict | bool | None = Field(57        default=None,58        description="Verifier support info: {sql: bool, code: bool} or legacy bool",59    )60 61    @field_validator("has_verifier", mode="before")62    @classmethod63    def _convert_bool_to_dict(cls, v: Any) -> dict | None:64        """Convert legacy bool format to new dict format."""65        if v is None:66            return None67        if isinstance(v, bool):68            # Legacy format: True means both modes available (conservative assumption)69            return {"sql": v, "code": v} if v else None70        return v71 72    num_tools: int | None = Field(73        default=None, description="Number of tools discovered"74    )75    tool_name: str | None = Field(default=None, description="Name of the tool called")76    tool_result: Any = Field(default=None, description="Result from the tool call")77    error: str | None = Field(default=None, description="Error message if any")78    warning: str | None = Field(default=None, description="Warning message if any")79    verify_result: dict | None = Field(80        default=None, description="Verifier output on episode end"81    )82    steps_taken: int | None = Field(83        default=None, description="Steps taken in this episode"84    )85    scenarios: list | None = Field(86        default=None, description="List of all scenarios (from __list_scenarios__)"87    )88    total: int | None = Field(default=None, description="Total number of scenarios")89    trajectory_path: str | None = Field(90        default=None, description="Path to saved trajectory JSON file"91    )92    session_dir: str | None = Field(93        default=None, description="Session directory path (when keep_session=True)"94    )95 96    def model_dump(self, **kwargs: Any) -> dict[str, Any]:97        kwargs.setdefault("exclude_none", True)98        return super().model_dump(**kwargs)99 100 101class AWMListToolsObservation(ListToolsObservation):102    """ListToolsObservation with AWM error field promoted to top level."""103 104    model_config = ConfigDict(extra="forbid")105 106    error: str | None = Field(default=None, description="Error message if any")107