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

plato/worlds/config.py

@@ -3,120 +3,286 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, Generic, TypeVar
+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 a single agent slot.
+    """Configuration for an agent.
 
     Attributes:
         image: Docker image URI for the agent
-        config: Agent-specific configuration (passed to the agent)
+        config: Agent-specific configuration passed to the agent
     """
 
     image: str
     config: dict[str, Any] = Field(default_factory=dict)
 
 
-class WorldConfig(BaseModel):
-    """World-specific configuration base class.
+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.
 
-    Subclass this to create typed configurations for your world:
+    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 environment fields.
 
-        class CodeWorldConfig(WorldConfig):
+    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, secrets, and envs:
+
+        class CodeWorldConfig(RunConfig):
+            # World-specific fields
             repository_url: str
-            checkout: str = "main"
             prompt: str
 
-    The schema will be auto-generated from the Pydantic model.
+            # Agents (typed)
+            coder: Annotated[AgentConfig, Agent(description="Coding agent")]
+
+            # 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
+        plato_session: Serialized Plato session for connecting to existing VM session
     """
 
+    session_id: str = ""
+    callback_url: str = ""
+    all_secrets: dict[str, str] = Field(default_factory=dict)  # All secrets (world + agent)
+
+    # 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 | None = None
+
     model_config = {"extra": "allow"}
 
-    def get(self, key: str, default: Any = None) -> Any:
-        """Get a config value by key (for backwards compatibility)."""
-        return getattr(self, key, default)
+    @classmethod
+    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, Env)):
+                    marker = meta
+                    break
+
+            result[field_name] = marker
+
+        return result
 
     @classmethod
     def get_json_schema(cls) -> dict:
-        """Get JSON schema for this config class."""
-        schema = cls.model_json_schema()
-        # Remove Pydantic-specific keys
-        schema.pop("title", None)
-        return schema
+        """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", {})
 
-# Type variable for typed config classes
-WorldConfigT = TypeVar("WorldConfigT", bound=WorldConfig)
+        world_properties = {}
+        agents = []
+        secrets = []
+        envs = []
 
+        # Skip runtime fields
+        runtime_fields = {"session_id", "callback_url", "all_secrets", "plato_session"}
 
-class RunConfig(BaseModel, Generic[WorldConfigT]):
-    """Full configuration for running a world.
+        for field_name, prop_schema in properties.items():
+            if field_name in runtime_fields:
+                continue
 
-    This is passed to the world runner and contains everything
-    needed to execute a world. Use with a type parameter for typed config access:
+            marker = annotations.get(field_name)
 
-        config: RunConfig[CodeWorldConfig] = ...
-        config.world_config.repository_url  # type-safe access
+            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
 
-    Attributes:
-        world_config: World-specific configuration (typed via generic)
-        agents: Mapping of slot name to agent configuration
-        secrets: Secret values (API keys, tokens, etc.)
-        session_id: Unique session identifier
-    """
+                envs.append(
+                    {
+                        "name": field_name,
+                        "description": marker.description,
+                        "required": marker.required,
+                        "default": default_value,
+                    }
+                )
+            else:
+                world_properties[field_name] = prop_schema
 
-    world_config: WorldConfigT = Field(default_factory=WorldConfig)  # type: ignore[assignment]
-    agents: dict[str, AgentConfig] = Field(default_factory=dict)
-    secrets: dict[str, str] = Field(default_factory=dict)
-    session_id: str = ""
-    callback_url: str = ""  # URL for agent to push logs back to Chronos
+        # 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
+        ]
 
-    def get_agent(self, slot_name: str) -> AgentConfig | None:
-        """Get agent config for a slot."""
-        return self.agents.get(slot_name)
+        return {
+            "properties": world_properties,
+            "required": required,
+            "agents": agents,
+            "secrets": secrets,
+            "envs": envs,
+        }
 
-    def get_secret(self, key: str, default: str | None = None) -> str | None:
-        """Get a secret value."""
-        return self.secrets.get(key, default)
+    def get_envs(self) -> list[EnvConfig]:
+        """Get all environment configurations from this config.
 
-    @classmethod
-    def from_file(
-        cls,
-        path: str | Path,
-        config_class: type[WorldConfig] | None = None,
-    ) -> RunConfig:
-        """Load run config from a JSON file.
-
-        Args:
-            path: Path to JSON config file
-            config_class: Optional WorldConfig subclass to parse world_config into.
-                         Defaults to WorldConfig if not provided.
+        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."""
         import json
 
         path = Path(path)
         with open(path) as f:
             data = json.load(f)
 
-        # Parse agents into AgentConfig objects
-        agents = {}
-        for slot_name, agent_data in data.get("agents", {}).items():
-            if isinstance(agent_data, dict):
-                agents[slot_name] = AgentConfig(**agent_data)
-            else:
-                agents[slot_name] = AgentConfig(image=str(agent_data))
-
-        # Use provided config class or default to WorldConfig
-        config_cls = config_class or WorldConfig
-        world_config = config_cls(**data.get("world_config", {}))
-
-        return cls(
-            world_config=world_config,  # type: ignore[arg-type]
-            agents=agents,
-            secrets=data.get("secrets", {}),
-            session_id=data.get("session_id", ""),
-            callback_url=data.get("callback_url", ""),
-        )
+        return cls._from_dict(data)
+
+    @classmethod
+    def _from_dict(cls, data: dict) -> RunConfig:
+        """Parse config from a dictionary, handling nested structures."""
+        annotations = cls.get_field_annotations()
+        parsed: dict[str, Any] = {}
+
+        # Handle world_config if present (for backwards compatibility)
+        world_config = data.pop("world_config", {})
+
+        # Handle agents dict -> individual agent fields
+        agents_dict = data.pop("agents", {})
+
+        # 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)
+
+        # Map agents dict to typed fields
+        for field_name, marker in annotations.items():
+            if isinstance(marker, Agent) and field_name in agents_dict:
+                agent_data = agents_dict[field_name]
+                if isinstance(agent_data, dict):
+                    parsed[field_name] = AgentConfig(**agent_data)
+                else:
+                    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)