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

plato/worlds/base.py

@@ -1,28 +1,39 @@
-"""Base world class for Plato worlds - OpenAI Gym-like interface."""
+"""Base world class for Plato worlds."""
 
 from __future__ import annotations
 
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, ClassVar
+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
 
 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 +56,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 +65,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"]
-
-            async def reset(self, config: RunConfig) -> Observation:
-                # Clone repo, setup workspace
-                return Observation(data={"workspace": "/workspace"})
+            description = "Run coding agents"
 
-            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 +132,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,25 +168,251 @@ 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}")
+
+    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
+
+        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:
+                self.logger.debug(f"{package_name} sim package not installed, skipping")
+            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.
+
+        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:
+                self.logger.debug(f"{package_name} sim package not installed, skipping instructions")
+            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 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}")
 
                 if result.done:
@@ -222,4 +420,17 @@ class BaseWorld(ABC):
 
         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")