ChilleD/agent_world_model_env

1"""2Agent World Model Environment Client.3 4Provides a client for connecting to an AWM Environment server.5AWMEnv extends MCPToolClient with AWM-specific helpers.6 7Example:8    >>> with AWMEnv(base_url="http://localhost:8000") as env:9    ...     # List all available scenarios10    ...     result = env.call_tool("__list_scenarios__")11    ...12    ...     # Start a scenario13    ...     env.reset(scenario="e_commerce_33", task_idx=0)14    ...15    ...     # Discover tools16    ...     tools = env.list_tools()17    ...     print([t.name for t in tools])18    ...19    ...     # Call tools20    ...     result = env.call_tool("search_products", query="headphones")21    ...22    ...     # End episode and get verification result23    ...     result = env.call_tool("done")24"""25 26from typing import Any27 28from openenv.core.client_types import StepResult29from openenv.core.mcp_client import MCPToolClient30 31from .models import AWMListToolsObservation, AWMObservation32 33 34class AWMEnv(MCPToolClient):35    """36    Client for the Agent World Model Environment.37    Inherits all functionality from MCPToolClient.38    AWM-specific reset parameters:39        scenario (str): Required. Name of the AWM scenario to load.40        task_idx (int): Optional. Index of the pre-defined task/verifier pair.41        task (str): Optional. Custom task description (no verifier).42        verifier_mode (str): "sql" or "code". Default: "sql".43        llm_base_url (str): LLM endpoint for sql verifier using any OpenAI compatible API service.44        llm_api_key (str): LLM API key.45        llm_model (str): LLM model name.46 47    Hidden tools (not in list_tools):48        done: End the episode and trigger verification.49        __list_scenarios__: List all 1,000 available scenarios.50 51    Accessing AWM-specific fields:52        After reset/step, access fields via observation attributes:53        >>> result = await env.reset(scenario="marketplace_1", task_idx=0)54        >>> print(result.observation.reward_type)  # "reset_ok"55        >>> print(result.observation.scenario)     # "marketplace_1"56        >>> print(result.observation.task)         # "Update my volunteer profile..."57    """58 59    def _parse_result(self, payload: dict[str, Any]) -> StepResult:60        """61        This override ensures AWM fields (reward_type, scenario, task, etc.)62        are available as observation attributes instead of being lost.63        """64        obs_data = payload.get("observation", {})65 66        # Check if this is a ListToolsObservation (has "tools" key)67        if "tools" in obs_data:68            observation = AWMListToolsObservation(**obs_data)69        else:70            observation = AWMObservation(**obs_data)71 72        return StepResult(73            observation=observation,74            reward=payload.get("reward"),75            done=payload.get("done", False),76        )77