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

plato/worlds/base.py

@@ -11,8 +11,14 @@ from pydantic import BaseModel, Field
 from plato.worlds.config import RunConfig
 
 if TYPE_CHECKING:
+    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
@@ -100,6 +106,7 @@ class BaseWorld(ABC, Generic[ConfigT]):
         self.logger = logging.getLogger(f"plato.worlds.{self.name}")
         self._step_count: int = 0
         self.plato_session = None
+        self._current_step_id: str | None = None
 
     @classmethod
     def get_config_class(cls) -> type[RunConfig]:
@@ -138,6 +145,7 @@ class BaseWorld(ABC, Generic[ConfigT]):
             "required": schema.get("required", []),
             "agents": schema.get("agents", []),
             "secrets": schema.get("secrets", []),
+            "envs": schema.get("envs", []),
         }
 
     @abstractmethod
@@ -166,6 +174,9 @@ class BaseWorld(ABC, Generic[ConfigT]):
         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
 
@@ -184,6 +195,163 @@ class BaseWorld(ABC, Generic[ConfigT]):
             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.
 
@@ -195,16 +363,56 @@ class BaseWorld(ABC, Generic[ConfigT]):
 
         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:
-            obs = await self.reset()
+            # 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}")
 
             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:
@@ -213,4 +421,16 @@ class BaseWorld(ABC, Generic[ConfigT]):
         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")

plato/worlds/config.py

@@ -7,8 +7,16 @@ from typing import Any
 
 from pydantic import BaseModel, Field
 
+from plato._generated.models import (
+    EnvFromArtifact,
+    EnvFromResource,
+    EnvFromSimulator,
+)
 from plato.v2.async_.session import SerializedSession
 
+# Union type for environment configurations
+EnvConfig = EnvFromArtifact | EnvFromSimulator | EnvFromResource
+
 
 class AgentConfig(BaseModel):
     """Configuration for an agent.
@@ -46,10 +54,28 @@ class Secret:
         self.required = required
 
 
+class Env:
+    """Annotation marker for environment fields.
+
+    Environments are VMs that run alongside the world's runtime.
+    They can be specified by artifact ID, simulator name, or resource config.
+
+    Usage:
+        gitea: Annotated[EnvConfig, Env(description="Git server")] = EnvFromArtifact(
+            artifact_id="abc123",
+            alias="gitea",
+        )
+    """
+
+    def __init__(self, description: str = "", required: bool = True):
+        self.description = description
+        self.required = required
+
+
 class RunConfig(BaseModel):
     """Base configuration for running a world.
 
-    Subclass this with your world-specific fields, agents, and secrets:
+    Subclass this with your world-specific fields, agents, secrets, and envs:
 
         class CodeWorldConfig(RunConfig):
             # World-specific fields
@@ -62,6 +88,12 @@ class RunConfig(BaseModel):
             # Secrets (typed)
             git_token: Annotated[str | None, Secret(description="GitHub token")] = None
 
+            # Environments (typed)
+            gitea: Annotated[EnvConfig, Env(description="Git server")] = EnvFromArtifact(
+                artifact_id="abc123",
+                alias="gitea",
+            )
+
     Attributes:
         session_id: Unique Chronos session identifier
         callback_url: Callback URL for status updates
@@ -74,21 +106,21 @@ class RunConfig(BaseModel):
 
     # Serialized Plato session for connecting to VM and sending heartbeats
     # This is the output of Session.dump() - used to restore session with Session.load()
-    plato_session: SerializedSession
+    plato_session: SerializedSession | None = None
 
     model_config = {"extra": "allow"}
 
     @classmethod
-    def get_field_annotations(cls) -> dict[str, Agent | Secret | None]:
-        """Get Agent/Secret annotations for each field."""
-        result: dict[str, Agent | Secret | None] = {}
+    def get_field_annotations(cls) -> dict[str, Agent | Secret | Env | None]:
+        """Get Agent/Secret/Env annotations for each field."""
+        result: dict[str, Agent | Secret | Env | None] = {}
 
         for field_name, field_info in cls.model_fields.items():
             marker = None
 
             # Pydantic stores Annotated metadata in field_info.metadata
             for meta in field_info.metadata:
-                if isinstance(meta, (Agent, Secret)):
+                if isinstance(meta, (Agent, Secret, Env)):
                     marker = meta
                     break
 
@@ -98,7 +130,7 @@ class RunConfig(BaseModel):
 
     @classmethod
     def get_json_schema(cls) -> dict:
-        """Get JSON schema with agents and secrets separated."""
+        """Get JSON schema with agents, secrets, and envs separated."""
         # Get full Pydantic schema
         full_schema = cls.model_json_schema()
         full_schema.pop("title", None)
@@ -110,6 +142,7 @@ class RunConfig(BaseModel):
         world_properties = {}
         agents = []
         secrets = []
+        envs = []
 
         # Skip runtime fields
         runtime_fields = {"session_id", "callback_url", "all_secrets", "plato_session"}
@@ -136,6 +169,26 @@ class RunConfig(BaseModel):
                         "required": marker.required,
                     }
                 )
+            elif isinstance(marker, Env):
+                # Get default value for this env field
+                field_info = cls.model_fields.get(field_name)
+                default_value = None
+                if field_info and field_info.default is not None:
+                    # Serialize the default EnvConfig to dict
+                    default_env = field_info.default
+                    if hasattr(default_env, "model_dump"):
+                        default_value = default_env.model_dump()
+                    elif isinstance(default_env, dict):
+                        default_value = default_env
+
+                envs.append(
+                    {
+                        "name": field_name,
+                        "description": marker.description,
+                        "required": marker.required,
+                        "default": default_value,
+                    }
+                )
             else:
                 world_properties[field_name] = prop_schema
 
@@ -149,8 +202,26 @@ class RunConfig(BaseModel):
             "required": required,
             "agents": agents,
             "secrets": secrets,
+            "envs": envs,
         }
 
+    def get_envs(self) -> list[EnvConfig]:
+        """Get all environment configurations from this config.
+
+        Returns:
+            List of EnvConfig objects (EnvFromArtifact, EnvFromSimulator, or EnvFromResource)
+        """
+        annotations = self.get_field_annotations()
+        envs: list[EnvConfig] = []
+
+        for field_name, marker in annotations.items():
+            if isinstance(marker, Env):
+                value = getattr(self, field_name, None)
+                if value is not None:
+                    envs.append(value)
+
+        return envs
+
     @classmethod
     def from_file(cls, path: str | Path) -> RunConfig:
         """Load config from a JSON file."""
@@ -177,6 +248,9 @@ class RunConfig(BaseModel):
         # Handle secrets dict -> individual secret fields
         secrets_dict = data.pop("secrets", {})
 
+        # Handle envs dict -> individual env fields
+        envs_dict = data.pop("envs", {})
+
         # Merge world_config into top-level
         parsed.update(world_config)
         parsed.update(data)
@@ -191,8 +265,24 @@ class RunConfig(BaseModel):
                     parsed[field_name] = AgentConfig(image=str(agent_data))
             elif isinstance(marker, Secret) and field_name in secrets_dict:
                 parsed[field_name] = secrets_dict[field_name]
+            elif isinstance(marker, Env) and field_name in envs_dict:
+                env_data = envs_dict[field_name]
+                if isinstance(env_data, dict):
+                    parsed[field_name] = _parse_env_config(env_data)
+                else:
+                    parsed[field_name] = env_data
 
         # Store all secrets for agent use
         parsed["all_secrets"] = secrets_dict
 
         return cls(**parsed)
+
+
+def _parse_env_config(data: dict) -> EnvConfig:
+    """Parse an env config dict into the appropriate type."""
+    if "artifact_id" in data:
+        return EnvFromArtifact(**data)
+    elif "sim_config" in data:
+        return EnvFromResource(**data)
+    else:
+        return EnvFromSimulator(**data)