plato-sdk-v2 2.0.64__py3-none-any.whl → 2.3.4__py3-none-any.whl

plato/worlds/base.py

@@ -1,28 +1,43 @@
-"""Base world class for Plato worlds - OpenAI Gym-like interface."""
+"""Base world class for Plato worlds."""
 
 from __future__ import annotations
 
 import logging
+import subprocess
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, ClassVar
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, get_args, get_origin
 
 from pydantic import BaseModel, Field
 
+from plato.worlds.config import RunConfig
+
 if TYPE_CHECKING:
-    from plato.worlds.config import RunConfig, WorldConfig
+    from plato.v2.async_.environment import Environment
+    from plato.v2.async_.session import Session
+
+from plato.agents.logging import init_logging as _init_chronos_logging
+from plato.agents.logging import log_event as _log_event
+from plato.agents.logging import reset_logging as _reset_chronos_logging
+from plato.agents.logging import span as _span
+from plato.agents.logging import upload_artifact as _upload_artifact
+from plato.agents.logging import upload_checkpoint as _upload_checkpoint
 
 logger = logging.getLogger(__name__)
 
 # Global registry of worlds
 _WORLD_REGISTRY: dict[str, type[BaseWorld]] = {}
 
+# Type variable for config
+ConfigT = TypeVar("ConfigT", bound=RunConfig)
+
 
 def register_world(name: str | None = None):
     """Decorator to register a world class.
 
     Usage:
         @register_world("code")
-        class CodeWorld(BaseWorld):
+        class CodeWorld(BaseWorld[CodeWorldConfig]):
             ...
     """
 
@@ -45,24 +60,8 @@ def get_world(name: str) -> type[BaseWorld] | None:
     return _WORLD_REGISTRY.get(name)
 
 
-class AgentSlot(BaseModel):
-    """Definition of an agent slot in a world."""
-
-    name: str
-    description: str = ""
-    required: bool = True
-
-
-class SecretSlot(BaseModel):
-    """Definition of a secret slot in a world."""
-
-    name: str
-    description: str = ""
-    required: bool = False
-
-
 class Observation(BaseModel):
-    """Observation returned from reset/step - what the agent sees."""
+    """Observation returned from reset/step."""
 
     data: dict[str, Any] = Field(default_factory=dict)
 
@@ -70,48 +69,60 @@ class Observation(BaseModel):
 
 
 class StepResult(BaseModel):
-    """Result of a step - OpenAI Gym style."""
+    """Result of a step."""
 
     observation: Observation
     done: bool = False
     info: dict[str, Any] = Field(default_factory=dict)
 
 
-class BaseWorld(ABC):
-    """Base class for Plato worlds - OpenAI Gym-like interface.
+class BaseWorld(ABC, Generic[ConfigT]):
+    """Base class for Plato worlds.
 
-    Subclasses implement:
-    - reset(): Setup the world, return initial observation
-    - step(): Execute one step, return StepResult
+    Subclass with a config type parameter for fully typed config access:
 
-    The base run() method handles the loop: reset -> step until done.
+        class CodeWorldConfig(RunConfig):
+            repository_url: str
+            prompt: str
+            coder: Annotated[AgentConfig, Agent(description="Coding agent")]
+            git_token: Annotated[str | None, Secret(description="GitHub token")] = None
 
-    Example:
         @register_world("code")
-        class CodeWorld(BaseWorld):
+        class CodeWorld(BaseWorld[CodeWorldConfig]):
             name = "code"
-            agents = ["coder"]
+            description = "Run coding agents"
 
-            async def reset(self, config: RunConfig) -> Observation:
-                # Clone repo, setup workspace
-                return Observation(data={"workspace": "/workspace"})
-
-            async def step(self) -> StepResult:
-                # Run agent, check if done
-                return StepResult(observation=obs, done=True)
+            async def reset(self) -> Observation:
+                url = self.config.repository_url  # typed as str
+                agent = self.config.coder          # typed as AgentConfig
+                token = self.config.git_token      # typed as str | None
     """
 
     # Class attributes
     name: ClassVar[str] = "base"
     description: ClassVar[str] = ""
-    agents: ClassVar[list[AgentSlot | str]] = []
-    secrets: ClassVar[list[SecretSlot | str]] = []
-    config_class: ClassVar[type[WorldConfig] | None] = None  # Override with typed config
+
+    # Instance attributes
+    config: ConfigT  # Typed via generic parameter
+    plato_session: Session | None = None  # Connected Plato session (if running on managed VM)
 
     def __init__(self) -> None:
         self.logger = logging.getLogger(f"plato.worlds.{self.name}")
-        self.config: RunConfig | None = None
         self._step_count: int = 0
+        self.plato_session = None
+        self._current_step_id: str | None = None
+
+    @classmethod
+    def get_config_class(cls) -> type[RunConfig]:
+        """Get the config class from the generic parameter."""
+        # Walk up the class hierarchy to find Generic base
+        for base in getattr(cls, "__orig_bases__", []):
+            origin = get_origin(base)
+            if origin is BaseWorld:
+                args = get_args(base)
+                if args and isinstance(args[0], type) and issubclass(args[0], RunConfig):
+                    return args[0]
+        return RunConfig
 
     @classmethod
     def get_version(cls) -> str:
@@ -125,62 +136,27 @@ class BaseWorld(ABC):
                 continue
         return "0.0.0"
 
-    @classmethod
-    def get_agent_slots(cls) -> list[AgentSlot]:
-        """Get the list of agent slots."""
-        return [AgentSlot(name=a) if isinstance(a, str) else a for a in cls.agents]
-
-    @classmethod
-    def get_secret_slots(cls) -> list[SecretSlot]:
-        """Get the list of secret slots."""
-        return [SecretSlot(name=s) if isinstance(s, str) else s for s in cls.secrets]
-
-    @classmethod
-    def get_config_schema(cls) -> dict:
-        """Get JSON schema for world_config.
-
-        If config_class is set, auto-generates schema from the Pydantic model.
-        Override this method to provide a custom schema.
-        """
-        if cls.config_class is not None:
-            return cls.config_class.get_json_schema()
-        return {"type": "object", "properties": {}, "required": []}
-
     @classmethod
     def get_schema(cls) -> dict:
-        """Get full schema including world_config, agents, and secrets."""
-        config_schema = cls.get_config_schema()
-        agent_slots = cls.get_agent_slots()
-        secret_slots = cls.get_secret_slots()
+        """Get full schema including world config, agents, and secrets."""
+        config_class = cls.get_config_class()
+        schema = config_class.get_json_schema()
 
         return {
             "$schema": "https://json-schema.org/draft/2020-12/schema",
             "type": "object",
-            "properties": config_schema.get("properties", {}),
-            "required": config_schema.get("required", []),
-            "agents": [
-                {
-                    "name": slot.name,
-                    "description": slot.description,
-                    "required": slot.required,
-                }
-                for slot in agent_slots
-            ],
-            "secrets": [
-                {
-                    "name": slot.name,
-                    "description": slot.description,
-                    "required": slot.required,
-                }
-                for slot in secret_slots
-            ],
+            "properties": schema.get("properties", {}),
+            "required": schema.get("required", []),
+            "agents": schema.get("agents", []),
+            "secrets": schema.get("secrets", []),
+            "envs": schema.get("envs", []),
         }
 
     @abstractmethod
-    async def reset(self, config: RunConfig) -> Observation:
+    async def reset(self) -> Observation:
         """Setup the world and return initial observation.
 
-        Called once at the start. Setup workspace, clone repos, etc.
+        Access configuration via self.config (fully typed).
         """
         pass
 
@@ -196,30 +172,539 @@ class BaseWorld(ABC):
         """Cleanup resources. Called after run completes."""
         pass
 
-    async def run(self, config: RunConfig) -> None:
+    async def _connect_plato_session(self) -> None:
+        """Connect to Plato session from config.
+
+        This is called automatically during run() to restore the session
+        and start sending heartbeats while the world runs.
+        """
+        if not self.config.plato_session:
+            return
+
+        try:
+            from plato.v2.async_.session import Session
+
+            self.logger.info("Restoring Plato session from serialized data")
+            self.plato_session = await Session.load(self.config.plato_session, start_heartbeat=True)
+            self.logger.info(f"Plato session {self.plato_session.session_id} restored, heartbeat started")
+        except Exception as e:
+            self.logger.warning(f"Failed to restore Plato session: {e}")
+
+    async def _disconnect_plato_session(self) -> None:
+        """Stop heartbeat for the Plato session (does not close the session)."""
+        if self.plato_session:
+            try:
+                await self.plato_session.stop_heartbeat()
+                self.logger.info("Plato session heartbeat stopped")
+            except Exception as e:
+                self.logger.warning(f"Error stopping Plato heartbeat: {e}")
+
+    async def _create_checkpoint(self) -> dict[str, str] | None:
+        """Create a checkpoint snapshot of all environments (excluding configured envs).
+
+        Uses snapshot_store for efficient chunk-based deduplication.
+
+        Returns:
+            Dict mapping environment alias to artifact_id, or None if no session connected.
+        """
+        if not self.plato_session:
+            self.logger.warning("Cannot create checkpoint: Plato session not connected")
+            return None
+
+        exclude_envs = set(self.config.checkpoint.exclude_envs)
+        envs_to_snapshot = [env for env in self.plato_session.envs if env.alias not in exclude_envs]
+
+        if not envs_to_snapshot:
+            self.logger.info("No environments to checkpoint (all excluded)")
+            return {}
+
+        self.logger.info(
+            f"Creating checkpoint for {len(envs_to_snapshot)} environment(s): {[e.alias for e in envs_to_snapshot]}"
+        )
+
+        results: dict[str, str] = {}
+        for env in envs_to_snapshot:
+            try:
+                result = await env.snapshot_store()
+                artifact_id = result.artifact_id
+                results[env.alias] = artifact_id
+
+                # Check for success/error fields (available after SDK regeneration)
+                success = getattr(result, "success", True)
+                error = getattr(result, "error", None)
+
+                if not success or error:
+                    self.logger.error(
+                        f"Checkpoint failed for '{env.alias}': {error or 'unknown error'} (job_id={env.job_id})"
+                    )
+                elif artifact_id:
+                    self.logger.info(f"Checkpoint created for '{env.alias}': {artifact_id}")
+                else:
+                    self.logger.warning(
+                        f"Checkpoint for '{env.alias}' returned empty artifact_id (job_id={env.job_id})"
+                    )
+            except Exception as e:
+                self.logger.error(f"Failed to checkpoint '{env.alias}': {e}")
+
+        return results
+
+    def _init_state_directory(self) -> None:
+        """Initialize the state directory as a git repository.
+
+        Creates the state directory if it doesn't exist and initializes it
+        as a git repository with an initial commit.
+        """
+        if not self.config.state.enabled:
+            return
+
+        state_path = Path(self.config.state.path)
+
+        # Create directory if it doesn't exist
+        if not state_path.exists():
+            state_path.mkdir(parents=True)
+            self.logger.info(f"Created state directory: {state_path}")
+
+        # Check if already a git repo
+        git_dir = state_path / ".git"
+        if git_dir.exists():
+            self.logger.info(f"State directory already initialized: {state_path}")
+            return
+
+        # Initialize git repo
+        try:
+            subprocess.run(
+                ["git", "init"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            # Create initial commit (even if empty)
+            subprocess.run(
+                ["git", "config", "user.email", "plato@plato.so"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            subprocess.run(
+                ["git", "config", "user.name", "Plato"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            # Add all files and create initial commit
+            subprocess.run(
+                ["git", "add", "-A"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            subprocess.run(
+                ["git", "commit", "--allow-empty", "-m", "Initial state"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            self.logger.info(f"Initialized git repo in state directory: {state_path}")
+        except subprocess.CalledProcessError as e:
+            self.logger.warning(f"Failed to initialize state git repo: {e.stderr}")
+
+    def _commit_state(self, message: str) -> bool:
+        """Commit current state directory changes.
+
+        Args:
+            message: Commit message
+
+        Returns:
+            True if commit was created (or no changes), False on error.
+        """
+        if not self.config.state.enabled:
+            return True
+
+        state_path = Path(self.config.state.path)
+        if not state_path.exists():
+            return True
+
+        try:
+            # Add all changes
+            subprocess.run(
+                ["git", "add", "-A"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            # Check if there are changes to commit
+            result = subprocess.run(
+                ["git", "status", "--porcelain"],
+                cwd=state_path,
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+            if not result.stdout.strip():
+                self.logger.debug("No state changes to commit")
+                return True
+
+            # Commit changes
+            subprocess.run(
+                ["git", "commit", "-m", message],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            self.logger.info(f"Committed state changes: {message}")
+            return True
+        except subprocess.CalledProcessError as e:
+            self.logger.warning(f"Failed to commit state: {e.stderr}")
+            return False
+
+    def _create_state_bundle(self) -> bytes | None:
+        """Create a git bundle of the state directory.
+
+        Returns:
+            Bundle bytes if successful, None otherwise.
+        """
+        if not self.config.state.enabled:
+            return None
+
+        state_path = Path(self.config.state.path)
+        if not state_path.exists():
+            return None
+
+        git_dir = state_path / ".git"
+        if not git_dir.exists():
+            self.logger.warning("State directory is not a git repository")
+            return None
+
+        try:
+            # Create bundle to stdout
+            result = subprocess.run(
+                ["git", "bundle", "create", "-", "--all"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            bundle_data = result.stdout
+            self.logger.info(f"Created state bundle: {len(bundle_data)} bytes")
+            return bundle_data
+        except subprocess.CalledProcessError as e:
+            self.logger.warning(f"Failed to create state bundle: {e.stderr}")
+            return None
+
+    async def _create_and_upload_checkpoint(self) -> dict[str, Any] | None:
+        """Create a full checkpoint including env snapshots and state bundle.
+
+        This method:
+        1. Commits any pending state changes
+        2. Creates env snapshots using snapshot_store
+        3. Creates and uploads state bundle as an artifact
+        4. Calls the checkpoint endpoint with all data
+
+        Returns:
+            Checkpoint result dict if successful, None otherwise.
+        """
+        # Commit state changes first
+        self._commit_state(f"Checkpoint at step {self._step_count}")
+
+        # Create env snapshots
+        env_snapshots = await self._create_checkpoint()
+        if env_snapshots is None:
+            env_snapshots = {}
+
+        # Create and upload state bundle
+        state_artifact_id: str | None = None
+        if self.config.state.enabled:
+            bundle_data = self._create_state_bundle()
+            if bundle_data:
+                result = await _upload_artifact(
+                    data=bundle_data,
+                    artifact_type="state",
+                    filename=f"state_step_{self._step_count}.bundle",
+                    extra={
+                        "step_number": self._step_count,
+                        "state_path": self.config.state.path,
+                    },
+                )
+                if result:
+                    state_artifact_id = result.get("artifact_id")
+                    self.logger.info(f"Uploaded state artifact: {state_artifact_id}")
+
+        # Upload checkpoint with all data
+        checkpoint_result = await _upload_checkpoint(
+            step_number=self._step_count,
+            env_snapshots=env_snapshots,
+            state_artifact_id=state_artifact_id,
+            extra={
+                "world_name": self.name,
+                "world_version": self.get_version(),
+            },
+        )
+
+        return checkpoint_result
+
+    def get_env(self, alias: str) -> Environment | None:
+        """Get an environment by alias.
+
+        Use this to access environments defined in the world config (e.g., gitea, localstack).
+
+        Args:
+            alias: The environment alias (e.g., "gitea", "localstack", "runtime")
+
+        Returns:
+            The Environment object or None if not found or session not connected.
+
+        Example:
+            gitea = self.get_env("gitea")
+            if gitea:
+                result = await gitea.execute("git status")
+        """
+        if not self.plato_session:
+            self.logger.warning("Cannot get env: Plato session not connected")
+            return None
+        return self.plato_session.get_env(alias)
+
+    @property
+    def envs(self) -> list[Environment]:
+        """Get all environments in the Plato session.
+
+        Returns:
+            List of Environment objects. Empty list if session not connected.
+        """
+        if not self.plato_session:
+            return []
+        return self.plato_session.envs
+
+    def get_sim_env_vars(self) -> dict[str, str]:
+        """Get environment variables from all configured sims.
+
+        Automatically discovers and loads env vars from sims like localstack, gitea, etc.
+        based on the environments configured in the world.
+
+        Returns:
+            Dict of environment variable name -> value
+
+        Raises:
+            ImportError: If a sim environment is configured but package is not installed.
+
+        Example:
+            env_vars = self.get_sim_env_vars()
+            # Returns: {"AWS_ENDPOINT_URL": "https://...", "GITEA_URL": "https://...", ...}
+        """
+        env_vars: dict[str, str] = {}
+
+        # Known sim packages and their env aliases
+        sim_packages = [
+            ("localstack", "localstack"),
+            ("gitea", "gitea"),
+        ]
+
+        for package_name, env_alias in sim_packages:
+            env = self.get_env(env_alias)
+            if not env:
+                continue
+
+            try:
+                # Dynamically import the sim package
+                sim_module = __import__(f"plato.sims.{package_name}", fromlist=[package_name])
+
+                # Get service URLs and env vars
+                service_urls = sim_module.get_service_urls(env.job_id)
+                sim_vars = sim_module.get_env_vars(service_urls)
+                env_vars.update(sim_vars)
+                self.logger.info(f"{package_name} env vars: {list(sim_vars.keys())}")
+            except ImportError:
+                raise ImportError(
+                    f"Environment '{env_alias}' is configured but 'plato.sims.{package_name}' "
+                    f"package is not installed.\n\n"
+                    f"Install sims packages:\n"
+                    f'  export INDEX_URL="https://__token__:${{PLATO_API_KEY}}@plato.so/api/v2/pypi/sims/simple/"\n'
+                    f"  uv pip install '.[sims]' --extra-index-url $INDEX_URL"
+                ) from None
+            except Exception as e:
+                self.logger.warning(f"Failed to get {package_name} env vars: {e}")
+
+        return env_vars
+
+    def get_sim_instructions(self) -> str:
+        """Get usage instructions from all configured sims.
+
+        Returns markdown-formatted instructions for using LocalStack, Gitea, etc.
+        based on the environments configured in the world.
+
+        Returns:
+            Markdown string with instructions, or empty string if no sims configured.
+
+        Raises:
+            ImportError: If a sim environment is configured but package is not installed.
+
+        Example:
+            instructions = self.get_sim_instructions()
+            # Returns markdown with LocalStack/Gitea setup instructions
+        """
+        instructions_parts: list[str] = []
+
+        # Known sim packages and their env aliases
+        sim_packages = [
+            ("localstack", "localstack"),
+            ("gitea", "gitea"),
+        ]
+
+        for package_name, env_alias in sim_packages:
+            env = self.get_env(env_alias)
+            if not env:
+                continue
+
+            try:
+                # Dynamically import the sim package
+                sim_module = __import__(f"plato.sims.{package_name}", fromlist=[package_name])
+
+                # Get instructions using the job_id
+                if hasattr(sim_module, "get_instructions_from_job"):
+                    instructions = sim_module.get_instructions_from_job(env.job_id)
+                    if instructions:
+                        instructions_parts.append(instructions)
+                        self.logger.info(f"Added {package_name} instructions to prompt")
+            except ImportError:
+                raise ImportError(
+                    f"Environment '{env_alias}' is configured but 'plato.sims.{package_name}' "
+                    f"package is not installed.\n\n"
+                    f"Install sims packages:\n"
+                    f'  export INDEX_URL="https://__token__:${{PLATO_API_KEY}}@plato.so/api/v2/pypi/sims/simple/"\n'
+                    f"  uv pip install '.[sims]' --extra-index-url $INDEX_URL"
+                ) from None
+            except Exception as e:
+                self.logger.warning(f"Failed to get {package_name} instructions: {e}")
+
+        if instructions_parts:
+            return "\n\n---\n\n".join(instructions_parts)
+        return ""
+
+    def format_instruction_with_sims(self, instruction: str) -> str:
+        """Format an instruction with sim context prepended.
+
+        Automatically adds available service instructions (LocalStack, Gitea, etc.)
+        before the main instruction.
+
+        Args:
+            instruction: The base instruction/task
+
+        Returns:
+            Formatted instruction with sim context, or original instruction if no sims.
+
+        Example:
+            formatted = self.format_instruction_with_sims("Fix the bug in main.py")
+            # Returns:
+            # ## Available Services
+            # The following services are available...
+            # [LocalStack instructions]
+            # ---
+            # ## Task
+            # Fix the bug in main.py
+        """
+        sim_instructions = self.get_sim_instructions()
+
+        if sim_instructions:
+            return f"""## Available Services
+
+The following services are available for your use:
+
+{sim_instructions}
+
+---
+
+## Task
+
+{instruction}"""
+        return instruction
+
+    async def run(self, config: ConfigT) -> None:
         """Run the world: reset -> step until done -> close.
 
-        This is the main entry point. Don't override this.
+        This is the main entry point. If plato_session_id is provided in config,
+        automatically connects to the Plato session to send heartbeats.
         """
         self.config = config
         self._step_count = 0
 
         self.logger.info(f"Starting world '{self.name}'")
 
+        # Initialize state directory (creates git repo if needed)
+        self._init_state_directory()
+
+        # Initialize the logging singleton for agents to use
+        if config.callback_url and config.session_id:
+            _init_chronos_logging(
+                callback_url=config.callback_url,
+                session_id=config.session_id,
+            )
+
+        # Connect to Plato session if configured (for heartbeats)
+        await self._connect_plato_session()
+
+        # Log session start
+        await _log_event(
+            span_type="session_start",
+            content=f"World '{self.name}' started",
+            source="world",
+            extra={"world_name": self.name, "world_version": self.get_version()},
+        )
+
         try:
-            # Reset (setup)
-            obs = await self.reset(config)
+            # Execute reset with automatic span tracking
+            async with _span("reset", span_type="reset", source="world") as reset_span:
+                reset_span.log(f"Resetting world '{self.name}'")
+                obs = await self.reset()
+                reset_span.set_extra({"observation": obs.model_dump() if hasattr(obs, "model_dump") else str(obs)})
             self.logger.info(f"World reset complete: {obs}")
 
-            # Step loop
             while True:
                 self._step_count += 1
-                result = await self.step()
+
+                # Execute step with automatic span tracking
+                # The span automatically sets itself as the current parent,
+                # so agent trajectories will nest under this step
+                async with _span(
+                    f"step_{self._step_count}",
+                    span_type="step",
+                    source="world",
+                ) as step_span:
+                    self._current_step_id = step_span.event_id
+                    step_span.log(f"Step {self._step_count} started")
+                    result = await self.step()
+                    step_span.set_extra(
+                        {
+                            "done": result.done,
+                            "observation": result.observation.model_dump()
+                            if hasattr(result.observation, "model_dump")
+                            else str(result.observation),
+                            "info": result.info,
+                        }
+                    )
+
                 self.logger.info(f"Step {self._step_count}: done={result.done}")
 
+                # Create checkpoint if enabled and interval matches
+                # Note: The checkpoint event is created by the callback endpoint,
+                # so we don't need a span wrapper here (would create duplicates)
+                if self.config.checkpoint.enabled and self._step_count % self.config.checkpoint.interval == 0:
+                    self.logger.info(f"Creating checkpoint after step {self._step_count}")
+                    await self._create_and_upload_checkpoint()
+
                 if result.done:
                     break
 
         finally:
             await self.close()
+            await self._disconnect_plato_session()
+
+            # Log session end
+            await _log_event(
+                span_type="session_end",
+                content=f"World '{self.name}' completed after {self._step_count} steps",
+                source="world",
+                extra={"total_steps": self._step_count},
+            )
+
+            # Reset the logging singleton
+            _reset_chronos_logging()
+
             self.logger.info(f"World '{self.name}' completed after {self._step_count} steps")