plato-sdk-v2 2.8.7__py3-none-any.whl → 2.9.0__py3-none-any.whl

plato/worlds/base.py

@@ -2,51 +2,43 @@
 
 from __future__ import annotations
 
+import asyncio
+import importlib.metadata
 import logging
 import os
 import subprocess
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, get_args, get_origin
-
-from pydantic import BaseModel, Field
+from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, get_args, get_origin
 
+from plato.agents.artifacts import upload_artifact as _upload_artifact_raw
+from plato.agents.otel import get_tracer, init_tracing, shutdown_tracing
+from plato.agents.runner import run_agent as _run_agent_raw
+from plato.v2.async_.session import Session
 from plato.worlds.config import RunConfig
+from plato.worlds.models import Observation, StepResult
+from plato.worlds.schema import get_world_schema
 
 if TYPE_CHECKING:
     from plato.v2.async_.environment import Environment
-    from plato.v2.async_.session import Session
-
-from plato.agents.artifacts import (
-    upload_artifact as _upload_artifact_raw,
-)
-from plato.agents.otel import (
-    get_tracer,
-    init_tracing,
-    shutdown_tracing,
-)
-from plato.agents.runner import run_agent as _run_agent_raw
 
 logger = logging.getLogger(__name__)
 
+# Global registry of worlds
+_WORLD_REGISTRY: dict[str, type[BaseWorld]] = {}
+
+# Type variable for config
+ConfigT = TypeVar("ConfigT", bound=RunConfig)
+
 
 def _get_plato_version() -> str:
     """Get the installed plato SDK version."""
     try:
-        from importlib.metadata import version
-
-        return version("plato")
+        return importlib.metadata.version("plato")
     except Exception:
         return "unknown"
 
 
-# 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.
 
@@ -75,22 +67,6 @@ def get_world(name: str) -> type[BaseWorld] | None:
     return _WORLD_REGISTRY.get(name)
 
 
-class Observation(BaseModel):
-    """Observation returned from reset/step."""
-
-    data: dict[str, Any] = Field(default_factory=dict)
-
-    model_config = {"extra": "allow"}
-
-
-class StepResult(BaseModel):
-    """Result of a step."""
-
-    observation: Observation
-    done: bool = False
-    info: dict[str, Any] = Field(default_factory=dict)
-
-
 class BaseWorld(ABC, Generic[ConfigT]):
     """Base class for Plato worlds.
 
@@ -144,8 +120,6 @@ class BaseWorld(ABC, Generic[ConfigT]):
     @classmethod
     def get_version(cls) -> str:
         """Get version from package metadata."""
-        import importlib.metadata
-
         for pkg_name in [cls.__module__.split(".")[0], f"plato-world-{cls.name}"]:
             try:
                 return importlib.metadata.version(pkg_name)
@@ -156,28 +130,7 @@ class BaseWorld(ABC, Generic[ConfigT]):
     @classmethod
     def get_schema(cls) -> dict:
         """Get full schema including world config, agents, secrets, and envs."""
-        config_class = cls.get_config_class()
-        schema = config_class.get_json_schema()
-
-        result = {
-            "$schema": "https://json-schema.org/draft/2020-12/schema",
-            "type": "object",
-            "properties": schema.get("properties", {}),
-            "required": schema.get("required", []),
-            "agents": schema.get("agents", []),
-            "secrets": schema.get("secrets", []),
-            "envs": schema.get("envs", []),
-        }
-
-        # Include env_list if present (for worlds with arbitrary environment lists)
-        if "env_list" in schema:
-            result["env_list"] = schema["env_list"]
-
-        # Include $defs if present (for nested type references)
-        if "$defs" in schema:
-            result["$defs"] = schema["$defs"]
-
-        return result
+        return get_world_schema(cls)
 
     @abstractmethod
     async def reset(self) -> Observation:
@@ -201,8 +154,6 @@ class BaseWorld(ABC, Generic[ConfigT]):
 
     async def _cleanup_agent_containers(self) -> None:
         """Stop any agent containers spawned by this world."""
-        import asyncio
-
         if not self._agent_containers:
             return
 
@@ -272,8 +223,6 @@ class BaseWorld(ABC, Generic[ConfigT]):
             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")
@@ -544,18 +493,18 @@ class BaseWorld(ABC, Generic[ConfigT]):
     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).
+        Use this to access environments defined in the world config.
 
         Args:
-            alias: The environment alias (e.g., "gitea", "localstack", "runtime")
+            alias: The environment alias (e.g., "runtime", "database")
 
         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")
+            env = self.get_env("runtime")
+            if env:
+                result = await env.execute("ls -la")
         """
         if not self.plato_session:
             self.logger.warning("Cannot get env: Plato session not connected")
@@ -576,8 +525,8 @@ class BaseWorld(ABC, Generic[ConfigT]):
     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.
+        Automatically discovers and loads env vars from sims based on the
+        environments configured in the world.
 
         Returns:
             Dict of environment variable name -> value
@@ -591,88 +540,54 @@ class BaseWorld(ABC, Generic[ConfigT]):
         """
         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
-
+        for env in self.envs:
             try:
-                # Dynamically import the sim package
-                sim_module = __import__(f"plato.sims.{package_name}", fromlist=[package_name])
+                sim_module = __import__(f"plato.sims.{env.alias}", fromlist=[env.alias])
+
+                if not hasattr(sim_module, "get_service_urls") or not hasattr(sim_module, "get_env_vars"):
+                    continue
 
-                # 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())}")
+                self.logger.info(f"{env.alias} 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
+                # Sim package not installed for this env - skip silently
+                continue
             except Exception as e:
-                self.logger.warning(f"Failed to get {package_name} env vars: {e}")
+                self.logger.warning(f"Failed to get {env.alias} 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-formatted instructions for using configured sims
+        based on the environments 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
+            # Returns markdown with sim 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
-
+        for env in self.envs:
             try:
-                # Dynamically import the sim package
-                sim_module = __import__(f"plato.sims.{package_name}", fromlist=[package_name])
+                sim_module = __import__(f"plato.sims.{env.alias}", fromlist=[env.alias])
 
-                # 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")
+                        self.logger.info(f"Added {env.alias} 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
+                # Sim package not installed for this env - skip silently
+                continue
             except Exception as e:
-                self.logger.warning(f"Failed to get {package_name} instructions: {e}")
+                self.logger.warning(f"Failed to get {env.alias} instructions: {e}")
 
         if instructions_parts:
             return "\n\n---\n\n".join(instructions_parts)
@@ -681,8 +596,7 @@ class BaseWorld(ABC, Generic[ConfigT]):
     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.
+        Automatically adds available service instructions before the main instruction.
 
         Args:
             instruction: The base instruction/task
@@ -695,7 +609,7 @@ class BaseWorld(ABC, Generic[ConfigT]):
             # Returns:
             # ## Available Services
             # The following services are available...
-            # [LocalStack instructions]
+            # [sim instructions]
             # ---
             # ## Task
             # Fix the bug in main.py

plato/worlds/config.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 from pathlib import Path
 from typing import Any
 
@@ -13,6 +14,8 @@ from plato._generated.models import (
     EnvFromSimulator,
 )
 from plato.v2.async_.session import SerializedSession
+from plato.worlds.markers import Agent, Env, EnvList, Secret
+from plato.worlds.schema import get_field_annotations, get_world_config_schema
 
 # Union type for environment configurations
 EnvConfig = EnvFromArtifact | EnvFromSimulator | EnvFromResource
@@ -30,66 +33,6 @@ class AgentConfig(BaseModel):
     config: dict[str, Any] = Field(default_factory=dict)
 
 
-class Agent:
-    """Annotation marker for agent fields.
-
-    Usage:
-        coder: Annotated[AgentConfig, Agent(description="Coding agent")]
-    """
-
-    def __init__(self, description: str = "", required: bool = True):
-        self.description = description
-        self.required = required
-
-
-class Secret:
-    """Annotation marker for secret fields.
-
-    Usage:
-        api_key: Annotated[str | None, Secret(description="API key")] = None
-    """
-
-    def __init__(self, description: str = "", required: bool = False):
-        self.description = description
-        self.required = required
-
-
-class Env:
-    """Annotation marker for single 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 EnvList:
-    """Annotation marker for a list of arbitrary environments.
-
-    Use this when the world accepts a dynamic list of environments
-    rather than named, fixed environment fields.
-
-    Usage:
-        envs: Annotated[
-            list[EnvFromSimulator | EnvFromArtifact | EnvFromResource],
-            EnvList(description="Environments to create for this task"),
-            Field(discriminator="type"),
-        ] = []
-    """
-
-    def __init__(self, description: str = ""):
-        self.description = description
-
-
 class StateConfig(BaseModel):
     """Configuration for world state persistence.
 
@@ -169,114 +112,12 @@ class RunConfig(BaseModel):
     @classmethod
     def get_field_annotations(cls) -> dict[str, Agent | Secret | Env | EnvList | None]:
         """Get Agent/Secret/Env/EnvList annotations for each field."""
-        result: dict[str, Agent | Secret | Env | EnvList | 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, Env, EnvList)):
-                    marker = meta
-                    break
-
-            result[field_name] = marker
-
-        return result
+        return get_field_annotations(cls)
 
     @classmethod
     def get_json_schema(cls) -> dict:
         """Get JSON schema with agents, secrets, and envs separated."""
-        # Get full Pydantic schema
-        full_schema = cls.model_json_schema()
-        full_schema.pop("title", None)
-
-        # Separate fields by annotation type
-        annotations = cls.get_field_annotations()
-        properties = full_schema.get("properties", {})
-
-        world_properties = {}
-        agents = []
-        secrets = []
-        envs = []
-        env_list_field: dict | None = None  # For EnvList marker (arbitrary envs)
-
-        # Skip runtime fields
-        runtime_fields = {"session_id", "otel_url", "upload_url", "plato_session", "checkpoint", "state"}
-
-        for field_name, prop_schema in properties.items():
-            if field_name in runtime_fields:
-                continue
-
-            marker = annotations.get(field_name)
-
-            if isinstance(marker, Agent):
-                agents.append(
-                    {
-                        "name": field_name,
-                        "description": marker.description,
-                        "required": marker.required,
-                    }
-                )
-            elif isinstance(marker, Secret):
-                secrets.append(
-                    {
-                        "name": field_name,
-                        "description": marker.description,
-                        "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,
-                    }
-                )
-            elif isinstance(marker, EnvList):
-                # Field marked as a list of arbitrary environments
-                env_list_field = {
-                    "name": field_name,
-                    "description": marker.description,
-                }
-            else:
-                world_properties[field_name] = prop_schema
-
-        # Compute required fields (excluding runtime and annotated fields)
-        required = [
-            r for r in full_schema.get("required", []) if r not in runtime_fields and annotations.get(r) is None
-        ]
-
-        result = {
-            "properties": world_properties,
-            "required": required,
-            "agents": agents,
-            "secrets": secrets,
-            "envs": envs,
-        }
-
-        # Include $defs if present (for nested type references)
-        if "$defs" in full_schema:
-            result["$defs"] = full_schema["$defs"]
-
-        # Add env_list if present (for worlds with arbitrary environment lists)
-        if env_list_field:
-            result["env_list"] = env_list_field
-
-        return result
+        return get_world_config_schema(cls)
 
     def get_envs(self) -> list[EnvConfig]:
         """Get all environment configurations from this config.
@@ -298,8 +139,6 @@ class RunConfig(BaseModel):
     @classmethod
     def from_file(cls, path: str | Path) -> RunConfig:
         """Load config from a JSON file."""
-        import json
-
         path = Path(path)
         with open(path) as f:
             data = json.load(f)

plato/worlds/markers.py

@@ -0,0 +1,63 @@
+"""Annotation markers for Plato world configurations."""
+
+from __future__ import annotations
+
+
+class Agent:
+    """Annotation marker for agent fields.
+
+    Usage:
+        coder: Annotated[AgentConfig, Agent(description="Coding agent")]
+    """
+
+    def __init__(self, description: str = "", required: bool = True):
+        self.description = description
+        self.required = required
+
+
+class Secret:
+    """Annotation marker for secret fields.
+
+    Usage:
+        api_key: Annotated[str | None, Secret(description="API key")] = None
+    """
+
+    def __init__(self, description: str = "", required: bool = False):
+        self.description = description
+        self.required = required
+
+
+class Env:
+    """Annotation marker for single 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:
+        database: Annotated[EnvConfig, Env(description="Database server")] = EnvFromArtifact(
+            artifact_id="abc123",
+            alias="database",
+        )
+    """
+
+    def __init__(self, description: str = "", required: bool = True):
+        self.description = description
+        self.required = required
+
+
+class EnvList:
+    """Annotation marker for a list of arbitrary environments.
+
+    Use this when the world accepts a dynamic list of environments
+    rather than named, fixed environment fields.
+
+    Usage:
+        envs: Annotated[
+            list[EnvFromSimulator | EnvFromArtifact | EnvFromResource],
+            EnvList(description="Environments to create for this task"),
+            Field(discriminator="type"),
+        ] = []
+    """
+
+    def __init__(self, description: str = ""):
+        self.description = description

plato/worlds/models.py

@@ -0,0 +1,23 @@
+"""Data models for Plato worlds."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class Observation(BaseModel):
+    """Observation returned from reset/step."""
+
+    data: dict[str, Any] = Field(default_factory=dict)
+
+    model_config = {"extra": "allow"}
+
+
+class StepResult(BaseModel):
+    """Result of a step."""
+
+    observation: Observation
+    done: bool = False
+    info: dict[str, Any] = Field(default_factory=dict)