plato-sdk-v2 2.0.50__py3-none-any.whl → 2.2.4__py3-none-any.whl

plato/world/base.py

@@ -1,267 +0,0 @@
-"""World base class - OpenAI Gym-like interface for Plato simulators."""
-
-from __future__ import annotations
-
-import logging
-from abc import ABC, abstractmethod
-from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, cast
-
-from plato.v2 import AsyncPlato, AsyncSession, EnvFromArtifact, EnvFromResource, EnvFromSimulator
-from plato.world.config import WorldConfig
-from plato.world.types import Observation, StepResult
-
-if TYPE_CHECKING:
-    from plato.storage import RolloutStorage
-
-logger = logging.getLogger(__name__)
-
-
-class World(ABC):
-    """OpenAI Gym-like environment backed by Plato simulators.
-
-    World manages the Plato session lifecycle and provides a gym-like interface
-    for agent interaction. Agents interact with simulators via the plato.sims SDK
-    using URLs exposed by the World.
-
-    Subclasses must implement:
-    - _on_reset(): Custom reset logic, returns initial Observation
-    - _on_step(): Custom step logic (e.g., run consumers), returns StepResult
-    - get_prompt(): Generate agent prompt for current state
-
-    Example:
-        class MerchantWorld(World):
-            async def _on_reset(self) -> Observation:
-                return Observation(step=0, date=self.config.start_date, data={"balance": 1000})
-
-            async def _on_step(self) -> StepResult:
-                # Run consumer simulation, calculate reward
-                return StepResult(observation=obs, reward=100.0)
-
-            def get_prompt(self) -> str:
-                return f"Day {self._current_step}: Manage your store..."
-
-        # Agent receives env vars:
-        # STORE_BASE_URL=http://...
-        # ACCOUNTING_BASE_URL=http://...
-        # And uses: from plato.sims import spree; client = spree.Client.create()
-    """
-
-    def __init__(self, config: WorldConfig, storage: RolloutStorage | None = None):
-        """Initialize World with configuration.
-
-        Args:
-            config: WorldConfig with environment specs and parameters
-            storage: Optional RolloutStorage for persisting data
-        """
-        self.config = config
-        self.storage = storage
-
-        # Plato session state
-        self._plato: AsyncPlato | None = None
-        self._session: AsyncSession | None = None
-        self._connect_urls: dict[str, str] = {}
-
-        # Simulation state
-        self._current_step: int = 0
-        self._current_date: datetime | None = None
-        self._done: bool = False
-
-    @property
-    def current_step(self) -> int:
-        """Get the current step number."""
-        return self._current_step
-
-    @property
-    def current_date(self) -> str:
-        """Get the current simulation date as ISO string."""
-        if self._current_date is None:
-            return self.config.start_date
-        return self._current_date.strftime("%Y-%m-%d")
-
-    @property
-    def is_done(self) -> bool:
-        """Check if the simulation is complete."""
-        return self._done
-
-    @property
-    def session(self) -> AsyncSession | None:
-        """Get the underlying Plato session."""
-        return self._session
-
-    async def reset(self) -> Observation:
-        """Reset world: create Plato session, return initial observation.
-
-        Creates a new Plato session with the configured environments and
-        waits for them to be ready. Returns the initial observation.
-
-        Returns:
-            Initial Observation from _on_reset()
-        """
-        logger.info(f"Resetting world '{self.config.name}'")
-
-        # Create Plato client and session
-        self._plato = AsyncPlato()
-        envs = [spec.to_plato_env() for spec in self.config.envs]
-
-        logger.info(f"Creating Plato session with {len(envs)} environments")
-        self._session = await self._plato.sessions.create(
-            envs=cast(list[EnvFromSimulator | EnvFromArtifact | EnvFromResource], envs),
-            timeout=self.config.session_timeout,
-        )
-
-        # Get connect URLs - these become env vars for agents
-        self._connect_urls = await self._session.get_connect_url()
-        logger.info(f"Got connect URLs: {list(self._connect_urls.keys())}")
-
-        # Initialize state
-        self._current_step = 0
-        self._current_date = datetime.fromisoformat(self.config.start_date)
-        self._done = False
-
-        # Call subclass hook
-        obs = await self._on_reset()
-
-        # Log to storage
-        if self.storage:
-            self.storage.log_reset(obs)
-
-        logger.info(f"World reset complete. Initial observation: step={obs.step}, date={obs.date}")
-        return obs
-
-    async def step(self) -> StepResult:
-        """Advance simulation by one step.
-
-        Increments the step counter, advances the simulation date,
-        and calls the subclass _on_step() hook for custom logic.
-
-        Returns:
-            StepResult with observation, reward, done, and info
-        """
-        if self._session is None:
-            raise RuntimeError("World not initialized. Call reset() first.")
-
-        self._current_step += 1
-
-        # Advance simulation date
-        await self._advance_date()
-
-        # Call subclass hook for custom step logic
-        result = await self._on_step()
-
-        # Update done status
-        if result.done:
-            self._done = True
-
-        # Log to storage
-        if self.storage:
-            self.storage.log_step(self._current_step, result)
-
-        logger.info(f"Step {self._current_step} complete. Reward: {result.reward:.2f}, Done: {result.done}")
-        return result
-
-    async def _advance_date(self) -> None:
-        """Advance the simulation date by one day.
-
-        Updates the date on all environments in the session.
-        """
-        if self._current_date is None:
-            self._current_date = datetime.fromisoformat(self.config.start_date)
-
-        self._current_date += timedelta(days=1)
-
-        if self._session:
-            try:
-                await self._session.set_date(self._current_date)
-                logger.debug(f"Advanced date to {self.current_date}")
-            except Exception as e:
-                logger.warning(f"Failed to set date: {e}")
-
-    async def close(self) -> None:
-        """Close Plato session and cleanup resources."""
-        logger.info(f"Closing world '{self.config.name}'")
-
-        if self._session:
-            try:
-                await self._session.close()
-            except Exception as e:
-                logger.warning(f"Error closing session: {e}")
-            self._session = None
-
-        if self._plato:
-            try:
-                await self._plato.close()
-            except Exception as e:
-                logger.warning(f"Error closing Plato client: {e}")
-            self._plato = None
-
-        self._connect_urls = {}
-
-    def get_env_vars(self) -> dict[str, str]:
-        """Get environment variables for agent.
-
-        Returns dict with simulator URLs and current simulation date:
-        {
-            "STORE_BASE_URL": "http://...",
-            "ACCOUNTING_BASE_URL": "http://...",
-            "SIMULATION_DATE": "2024-01-15",
-        }
-
-        Agent can then: from plato.sims import spree; client = spree.Client.create()
-        """
-        env_vars = {f"{alias.upper()}_BASE_URL": url for alias, url in self._connect_urls.items()}
-        env_vars["SIMULATION_DATE"] = self.current_date
-        env_vars["SIMULATION_STEP"] = str(self._current_step)
-        return env_vars
-
-    def get_connect_urls(self) -> dict[str, str]:
-        """Get raw connect URLs by alias.
-
-        Returns:
-            Dict mapping alias to connect URL
-        """
-        return self._connect_urls.copy()
-
-    @abstractmethod
-    async def _on_reset(self) -> Observation:
-        """Subclass hook: custom reset logic.
-
-        Called after Plato session is created and ready.
-        Should return the initial observation.
-
-        Returns:
-            Initial Observation
-        """
-        pass
-
-    @abstractmethod
-    async def _on_step(self) -> StepResult:
-        """Subclass hook: custom step logic.
-
-        Called after date is advanced. Should implement any custom
-        step logic (e.g., run consumer simulation) and return the result.
-
-        Returns:
-            StepResult with observation, reward, done, info
-        """
-        pass
-
-    @abstractmethod
-    def get_prompt(self) -> str:
-        """Generate agent prompt for current state.
-
-        Should return a prompt string that describes what the agent
-        should do in this step.
-
-        Returns:
-            Prompt string for the agent
-        """
-        pass
-
-    async def __aenter__(self) -> World:
-        """Async context manager entry."""
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Async context manager exit."""
-        await self.close()

plato/world/config.py

@@ -1,139 +0,0 @@
-"""World configuration - WorldConfig and EnvSpec."""
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Any
-
-import yaml
-from pydantic import BaseModel, Field, model_validator
-
-from plato.v2 import Env, EnvFromArtifact, EnvFromSimulator
-
-
-class EnvSpec(BaseModel):
-    """Specification for a single environment in the world.
-
-    Environments can be specified either by:
-    - simulator + dataset: Create a fresh simulator instance
-    - artifact_id: Load from a saved artifact
-
-    Attributes:
-        alias: Unique identifier for this environment (e.g., "store", "accounting")
-        simulator: Simulator type (e.g., "spree", "firefly", "espocrm")
-        artifact_id: Optional artifact ID to load from
-        dataset: Dataset to use when creating fresh instance (default: "blank")
-    """
-
-    alias: str
-    simulator: str | None = None
-    artifact_id: str | None = None
-    dataset: str | None = None
-
-    @model_validator(mode="after")
-    def validate_env_source(self) -> EnvSpec:
-        """Ensure either simulator or artifact_id is provided."""
-        if not self.artifact_id and not self.simulator:
-            raise ValueError(f"EnvSpec '{self.alias}' must have either simulator or artifact_id")
-        return self
-
-    def to_plato_env(self) -> EnvFromArtifact | EnvFromSimulator:
-        """Convert to Plato Env object for session creation.
-
-        Returns:
-            Env object suitable for AsyncPlato.sessions.create()
-        """
-        if self.artifact_id:
-            return Env.artifact(self.artifact_id, alias=self.alias)
-        # validator ensures simulator is set if artifact_id is not
-        assert self.simulator is not None
-        return Env.simulator(self.simulator, alias=self.alias, dataset=self.dataset)
-
-
-class WorldConfig(BaseModel):
-    """Configuration for a World.
-
-    All configs can be loaded from YAML files.
-
-    Example YAML format:
-        ```yaml
-        name: "merchant-sim"
-        num_steps: 30
-        start_date: "2024-01-01"
-        envs:
-          - alias: store
-            simulator: spree
-          - alias: accounting
-            artifact_id: "abc-123"
-        ```
-
-    Attributes:
-        name: Human-readable name for this world
-        envs: List of environment specifications
-        num_steps: Number of simulation steps (e.g., days)
-        start_date: Simulation start date (ISO format)
-        seed: Random seed for reproducibility
-    """
-
-    name: str = "world"
-    envs: list[EnvSpec] = Field(default_factory=list)
-    num_steps: int = 30
-    start_date: str = "2024-01-01"
-    seed: int = 42
-    session_timeout: int = Field(default=300, description="Timeout for session creation in seconds")
-
-    @classmethod
-    def from_yaml(cls, path: str | Path) -> WorldConfig:
-        """Load configuration from a YAML file.
-
-        Args:
-            path: Path to YAML config file
-
-        Returns:
-            Loaded WorldConfig instance
-
-        Raises:
-            FileNotFoundError: If config file doesn't exist
-        """
-        path = Path(path)
-        if not path.exists():
-            raise FileNotFoundError(f"Config file not found: {path}")
-
-        with open(path) as f:
-            data = yaml.safe_load(f)
-
-        return cls(**data)
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> WorldConfig:
-        """Create configuration from a dictionary.
-
-        Args:
-            data: Configuration dictionary
-
-        Returns:
-            WorldConfig instance
-        """
-        return cls(**data)
-
-    def get_env(self, alias: str) -> EnvSpec | None:
-        """Get environment spec by alias.
-
-        Args:
-            alias: Environment alias to find
-
-        Returns:
-            EnvSpec if found, None otherwise
-        """
-        for env in self.envs:
-            if env.alias == alias:
-                return env
-        return None
-
-    def get_env_aliases(self) -> list[str]:
-        """Get list of all environment aliases.
-
-        Returns:
-            List of alias strings
-        """
-        return [env.alias for env in self.envs]

plato/world/types.py

@@ -1,47 +0,0 @@
-"""World types - Observation and StepResult."""
-
-from __future__ import annotations
-
-from typing import Any
-
-from pydantic import BaseModel, Field
-
-
-class Observation(BaseModel):
-    """Observation from the world after a step.
-
-    Contains the current state information that the agent observes.
-
-    Attributes:
-        step: Current step number (0 for initial reset)
-        date: Current simulation date (ISO format)
-        data: Custom observation data (world-specific)
-    """
-
-    step: int
-    date: str
-    data: dict[str, Any] = Field(default_factory=dict)
-
-    model_config = {"extra": "allow"}
-
-
-class StepResult(BaseModel):
-    """Result of a world step.
-
-    Follows OpenAI Gym convention with observation, reward, done, info.
-
-    Attributes:
-        observation: Current state/observation after the step
-        reward: Numeric reward for this step (e.g., profit, score)
-        done: Whether the simulation has ended
-        truncated: Whether step was cut short (e.g., timeout)
-        info: Additional information about the step
-    """
-
-    observation: Observation
-    reward: float = 0.0
-    done: bool = False
-    truncated: bool = False
-    info: dict[str, Any] = Field(default_factory=dict)
-
-    model_config = {"arbitrary_types_allowed": True}