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

plato/agents/base.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import importlib.metadata
 import json
 import logging
 from abc import ABC, abstractmethod
@@ -9,6 +10,7 @@ from pathlib import Path
 from typing import Any, ClassVar, Generic, TypeVar, get_args, get_origin
 
 from plato.agents.config import AgentConfig
+from plato.agents.schema import get_agent_schema
 
 logger = logging.getLogger(__name__)
 
@@ -91,8 +93,6 @@ class BaseAgent(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-agent-{cls.name}"]:
             try:
                 return importlib.metadata.version(pkg_name)
@@ -103,13 +103,7 @@ class BaseAgent(ABC, Generic[ConfigT]):
     @classmethod
     def get_schema(cls) -> dict:
         """Get full schema for the agent including config and build schemas."""
-        from plato.agents.build import BuildConfig
-
-        config_class = cls.get_config_class()
-        return {
-            "config": config_class.get_json_schema(),
-            "build": BuildConfig.get_json_schema(),
-        }
+        return get_agent_schema(cls)
 
     @abstractmethod
     async def run(self, instruction: str) -> None:

plato/agents/config.py

@@ -17,25 +17,15 @@ Example:
 
 from __future__ import annotations
 
+import base64
 import json
+import os
 from typing import Any
 
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
-
-class Secret:
-    """Annotation marker for secret fields.
-
-    Fields annotated with Secret are automatically loaded from environment variables.
-    The env var name is the uppercase version of the field name (e.g., api_key -> API_KEY).
-
-    Usage:
-        api_key: Annotated[str, Secret(description="API key")]
-    """
-
-    def __init__(self, description: str = "", required: bool = False):
-        self.description = description
-        self.required = required
+from plato.agents.markers import Secret
+from plato.agents.schema import get_agent_config_schema, get_field_secrets
 
 
 class AgentConfig(BaseSettings):
@@ -68,56 +58,12 @@ class AgentConfig(BaseSettings):
     @classmethod
     def get_field_secrets(cls) -> dict[str, Secret]:
         """Get Secret annotations for each field."""
-        result: dict[str, Secret] = {}
-
-        for field_name, field_info in cls.model_fields.items():
-            for meta in field_info.metadata:
-                if isinstance(meta, Secret):
-                    result[field_name] = meta
-                    break
-
-        return result
+        return get_field_secrets(cls)
 
     @classmethod
     def get_json_schema(cls) -> dict:
         """Get JSON schema with secrets separated."""
-        full_schema = cls.model_json_schema()
-        full_schema.pop("title", None)
-
-        secrets_map = cls.get_field_secrets()
-        properties = full_schema.get("properties", {})
-
-        config_properties = {}
-        secrets = []
-
-        # Skip internal fields
-        internal_fields = {"logs_dir"}
-
-        for field_name, prop_schema in properties.items():
-            if field_name in internal_fields:
-                continue
-
-            if field_name in secrets_map:
-                secret = secrets_map[field_name]
-                secrets.append(
-                    {
-                        "name": field_name,
-                        "description": secret.description,
-                        "required": secret.required,
-                    }
-                )
-            else:
-                config_properties[field_name] = prop_schema
-
-        required = [r for r in full_schema.get("required", []) if r not in internal_fields and r not in secrets_map]
-
-        return {
-            "$schema": "https://json-schema.org/draft/2020-12/schema",
-            "type": "object",
-            "properties": config_properties,
-            "required": required,
-            "secrets": secrets,
-        }
+        return get_agent_config_schema(cls)
 
     def get_secrets_dict(self) -> dict[str, str]:
         """Extract secret values as a dict for environment variables."""
@@ -137,7 +83,7 @@ class AgentConfig(BaseSettings):
         internal_fields = {"logs_dir"}
 
         result: dict[str, Any] = {}
-        for field_name in self.model_fields:
+        for field_name in self.__class__.model_fields:
             if field_name not in secrets_map and field_name not in internal_fields:
                 value = getattr(self, field_name, None)
                 if value is not None:
@@ -152,9 +98,6 @@ class AgentConfig(BaseSettings):
         The runner passes config as base64-encoded JSON in the
         AGENT_CONFIG_B64 environment variable.
         """
-        import base64
-        import os
-
         config_b64 = os.environ.get("AGENT_CONFIG_B64")
         if not config_b64:
             raise ValueError("AGENT_CONFIG_B64 environment variable not set")

plato/agents/markers.py

@@ -0,0 +1,18 @@
+"""Annotation markers for Plato agent configurations."""
+
+from __future__ import annotations
+
+
+class Secret:
+    """Annotation marker for secret fields.
+
+    Fields annotated with Secret are automatically loaded from environment variables.
+    The env var name is the uppercase version of the field name (e.g., api_key -> API_KEY).
+
+    Usage:
+        api_key: Annotated[str, Secret(description="API key")]
+    """
+
+    def __init__(self, description: str = "", required: bool = False):
+        self.description = description
+        self.required = required

plato/agents/otel.py

@@ -26,11 +26,19 @@ Usage:
 from __future__ import annotations
 
 import logging
+import os
 
+from opentelemetry import context as context_api
 from opentelemetry import trace
-from opentelemetry.trace import Tracer
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+    OTLPSpanExporter,
+)
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, Tracer
 
-_module_logger = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 # Global state
 _tracer_provider = None
@@ -39,10 +47,7 @@ _log_handler = None
 
 
 class OTelSpanLogHandler(logging.Handler):
-    """Logging handler that creates OTel spans for log messages.
-
-    Converts Python log records to OTel spans with log attributes.
-    """
+    """Logging handler that creates OTel spans for log messages."""
 
     def __init__(self, tracer: Tracer, level: int = logging.INFO):
         super().__init__(level)
@@ -51,9 +56,6 @@ class OTelSpanLogHandler(logging.Handler):
     def emit(self, record: logging.LogRecord) -> None:
         """Emit a log record as an OTel span."""
         try:
-            # Debug: print that we're emitting a log span
-            print(f"[OTel] Emitting log span: {record.name} - {record.getMessage()[:100]}")
-            # Create a span for the log message
             with self.tracer.start_as_current_span(f"log.{record.levelname.lower()}") as span:
                 span.set_attribute("log.level", record.levelname)
                 span.set_attribute("log.message", record.getMessage())
@@ -66,12 +68,9 @@ class OTelSpanLogHandler(logging.Handler):
                 if record.lineno:
                     span.set_attribute("log.lineno", record.lineno)
 
-                # Mark errors
                 if record.levelno >= logging.ERROR:
                     span.set_attribute("error", True)
-
         except Exception:
-            # Don't let logging errors crash the application
             pass
 
 
@@ -94,20 +93,10 @@ def init_tracing(
     global _tracer_provider, _initialized, _log_handler
 
     if _initialized:
-        _module_logger.debug("Tracing already initialized")
+        logger.debug("Tracing already initialized")
         return
 
     try:
-        from opentelemetry import context as context_api
-        from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
-            OTLPSpanExporter,
-        )
-        from opentelemetry.sdk.resources import Resource
-        from opentelemetry.sdk.trace import TracerProvider
-        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
-        from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
-
-        # Create resource with session ID
         resource = Resource.create(
             {
                 "service.name": service_name,
@@ -115,53 +104,39 @@ def init_tracing(
             }
         )
 
-        # Create tracer provider
         _tracer_provider = TracerProvider(resource=resource)
 
-        # Add OTLP exporter pointing to Chronos (use SimpleSpanProcessor for immediate export)
         otlp_exporter = OTLPSpanExporter(endpoint=f"{otlp_endpoint.rstrip('/')}/v1/traces")
         _tracer_provider.add_span_processor(SimpleSpanProcessor(otlp_exporter))
 
-        # Set as global tracer provider
         trace.set_tracer_provider(_tracer_provider)
 
-        # If parent context is provided, set it as the current context
-        # This allows new spans to automatically link to the parent
         if parent_trace_id and parent_span_id:
             parent_context = SpanContext(
                 trace_id=int(parent_trace_id, 16),
                 span_id=int(parent_span_id, 16),
                 is_remote=True,
-                trace_flags=TraceFlags(0x01),  # Sampled
+                trace_flags=TraceFlags(0x01),
             )
             parent_span = NonRecordingSpan(parent_context)
             ctx = trace.set_span_in_context(parent_span)
             context_api.attach(ctx)
-            print(f"[OTel] Using parent context: trace_id={parent_trace_id}, span_id={parent_span_id}")
+            logger.debug(f"Using parent context: trace_id={parent_trace_id}, span_id={parent_span_id}")
 
-        # Add OTel logging handler to capture logs from plato SDK
         tracer = trace.get_tracer(service_name)
         _log_handler = OTelSpanLogHandler(tracer, level=logging.INFO)
 
-        # Add handler to plato loggers (worlds and agents)
-        # Set level to INFO to ensure logs propagate from child loggers
         plato_logger = logging.getLogger("plato")
         plato_logger.setLevel(logging.INFO)
         plato_logger.addHandler(_log_handler)
-        print(
-            f"[OTel] Added log handler to 'plato' logger (level={plato_logger.level}, handlers={len(plato_logger.handlers)})"
-        )
 
         _initialized = True
-
-        print(f"[OTel] Tracing initialized: service={service_name}, session={session_id}, endpoint={otlp_endpoint}")
+        logger.debug(f"Tracing initialized: service={service_name}, session={session_id}")
 
     except ImportError as e:
-        print(f"[OTel] OpenTelemetry SDK not installed: {e}")
-        _module_logger.warning(f"OpenTelemetry SDK not installed: {e}")
+        logger.warning(f"OpenTelemetry SDK not installed: {e}")
     except Exception as e:
-        print(f"[OTel] Failed to initialize tracing: {e}")
-        _module_logger.error(f"Failed to initialize tracing: {e}")
+        logger.error(f"Failed to initialize tracing: {e}")
 
 
 def shutdown_tracing(timeout_millis: int = 30000) -> None:
@@ -172,7 +147,6 @@ def shutdown_tracing(timeout_millis: int = 30000) -> None:
     """
     global _tracer_provider, _initialized, _log_handler
 
-    # Remove log handler
     if _log_handler:
         try:
             plato_logger = logging.getLogger("plato")
@@ -183,19 +157,13 @@ def shutdown_tracing(timeout_millis: int = 30000) -> None:
 
     if _tracer_provider:
         try:
-            # Force flush all pending spans before shutdown
-            print(f"[OTel] Flushing spans (timeout={timeout_millis}ms)...")
             flush_success = _tracer_provider.force_flush(timeout_millis=timeout_millis)
-            if flush_success:
-                print("[OTel] Span flush completed successfully")
-            else:
-                print("[OTel] Span flush timed out or failed")
-
+            if not flush_success:
+                logger.warning("Span flush timed out")
             _tracer_provider.shutdown()
-            print("[OTel] Tracing shutdown complete")
+            logger.debug("Tracing shutdown complete")
         except Exception as e:
-            print(f"[OTel] Error shutting down tracer: {e}")
-            _module_logger.warning(f"Error shutting down tracer: {e}")
+            logger.warning(f"Error shutting down tracer: {e}")
 
     _tracer_provider = None
     _initialized = False
@@ -235,21 +203,15 @@ def instrument(service_name: str = "plato-agent") -> Tracer:
     Returns:
         OpenTelemetry Tracer
     """
-    import os
-
     otel_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
     session_id = os.environ.get("SESSION_ID", "local")
     parent_trace_id = os.environ.get("OTEL_TRACE_ID")
     parent_span_id = os.environ.get("OTEL_PARENT_SPAN_ID")
 
-    print(f"[OTel] instrument() called: service={service_name}, endpoint={otel_endpoint}, session={session_id}")
-
     if not otel_endpoint:
-        # Return default tracer (no-op if no provider configured)
-        print("[OTel] No OTEL_EXPORTER_OTLP_ENDPOINT set, returning no-op tracer")
+        logger.debug("No OTEL_EXPORTER_OTLP_ENDPOINT set, using no-op tracer")
         return trace.get_tracer(service_name)
 
-    # Initialize tracing with parent context if provided
     init_tracing(
         service_name=service_name,
         session_id=session_id,

plato/agents/schema.py

@@ -0,0 +1,96 @@
+"""JSON schema generation for Plato agent configurations."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from plato.agents.markers import Secret
+
+if TYPE_CHECKING:
+    from plato.agents.config import AgentConfig
+
+
+def get_field_secrets(config_cls: type[AgentConfig]) -> dict[str, Secret]:
+    """Get Secret annotations for each field.
+
+    Args:
+        config_cls: The AgentConfig subclass to inspect
+
+    Returns:
+        Dict mapping field name to Secret annotation
+    """
+    result: dict[str, Secret] = {}
+
+    for field_name, field_info in config_cls.model_fields.items():
+        for meta in field_info.metadata:
+            if isinstance(meta, Secret):
+                result[field_name] = meta
+                break
+
+    return result
+
+
+def get_agent_config_schema(config_cls: type[AgentConfig]) -> dict[str, Any]:
+    """Get JSON schema for an agent config with secrets separated.
+
+    Args:
+        config_cls: The AgentConfig subclass to generate schema for
+
+    Returns:
+        JSON schema dict with properties, required, and secrets fields
+    """
+    full_schema = config_cls.model_json_schema()
+    full_schema.pop("title", None)
+
+    secrets_map = get_field_secrets(config_cls)
+    properties = full_schema.get("properties", {})
+
+    config_properties: dict[str, Any] = {}
+    secrets: list[dict[str, Any]] = []
+
+    # Skip internal fields
+    internal_fields = {"logs_dir"}
+
+    for field_name, prop_schema in properties.items():
+        if field_name in internal_fields:
+            continue
+
+        if field_name in secrets_map:
+            secret = secrets_map[field_name]
+            secrets.append(
+                {
+                    "name": field_name,
+                    "description": secret.description,
+                    "required": secret.required,
+                }
+            )
+        else:
+            config_properties[field_name] = prop_schema
+
+    required = [r for r in full_schema.get("required", []) if r not in internal_fields and r not in secrets_map]
+
+    return {
+        "$schema": "https://json-schema.org/draft/2020-12/schema",
+        "type": "object",
+        "properties": config_properties,
+        "required": required,
+        "secrets": secrets,
+    }
+
+
+def get_agent_schema(agent_cls: type) -> dict[str, Any]:
+    """Get full schema for an agent including config and build schemas.
+
+    Args:
+        agent_cls: The BaseAgent subclass to generate schema for
+
+    Returns:
+        Dict with config and build schema sections
+    """
+    from plato.agents.build import BuildConfig
+
+    config_class = agent_cls.get_config_class()
+    return {
+        "config": get_agent_config_schema(config_class),
+        "build": BuildConfig.get_json_schema(),
+    }

plato/cli/chronos.py

@@ -6,11 +6,13 @@ import asyncio
 import json
 import logging
 import os
+import platform as plat
 import subprocess
 import tempfile
 from pathlib import Path
 from typing import Annotated
 
+import httpx
 import typer
 
 from plato.cli.utils import console
@@ -64,8 +66,6 @@ def launch(
     The config file should contain world.package (required) and optionally world.config,
     runtime.artifact_id, and tags.
     """
-    import httpx
-
     # Set defaults
     if not chronos_url:
         chronos_url = "https://chronos.plato.so"
@@ -126,80 +126,6 @@ def launch(
         raise typer.Exit(1)
 
 
-@chronos_app.command()
-def example(
-    world: str = typer.Argument(
-        "structured-execution",
-        help="World to generate example config for",
-    ),
-    output: Path = typer.Option(
-        None,
-        "--output",
-        "-o",
-        help="Output file path (prints to stdout if not specified)",
-    ),
-):
-    """Generate an example job config file.
-
-    Creates a sample JSON configuration for launching Chronos jobs, which can be
-    customized for your use case.
-
-    Arguments:
-        world: World type to generate example for (default: "structured-execution")
-
-    Options:
-        -o, --output: Output file path. If not specified, prints to stdout.
-
-    Available worlds: structured-execution, code-world
-    """
-    examples = {
-        "structured-execution": {
-            "world_package": "plato-world-structured-execution",
-            "world_version": "latest",
-            "world_config": {
-                "sim_name": "my-sim",
-                "github_url": "https://github.com/example/repo",
-                "max_attempts": 3,
-                "use_backtrack": True,
-                "skill_runner": {
-                    "image": "claude-code:2.1.5",
-                    "config": {"model_name": "anthropic/claude-sonnet-4-20250514", "max_turns": 100},
-                },
-                "plato_api_key": "pk_xxx",
-                "anthropic_api_key": "sk-ant-xxx",
-            },
-            "_comment": "Agents and secrets are embedded directly in world_config",
-        },
-        "code-world": {
-            "world_package": "plato-world-code",
-            "world_config": {
-                "task": "Fix the bug in src/main.py",
-                "repo_url": "https://github.com/example/repo",
-                "coder": {
-                    "image": "claude-code:latest",
-                    "config": {"model_name": "anthropic/claude-sonnet-4-20250514"},
-                },
-            },
-            "_comment": "world_version is optional - uses latest if not specified",
-        },
-    }
-
-    if world not in examples:
-        console.print(f"[red]❌ Unknown world: {world}[/red]")
-        console.print(f"Available examples: {list(examples.keys())}")
-        raise typer.Exit(1)
-
-    example_config = examples[world]
-    json_output = json.dumps(example_config, indent=2)
-
-    if output:
-        with open(output, "w") as f:
-            f.write(json_output)
-        console.print(f"[green]✅ Example config written to {output}[/green]")
-    else:
-        console.print(json_output)
-
-
 def _get_world_runner_dockerfile() -> Path:
     """Get the path to the world runner Dockerfile template."""
     return Path(__file__).parent.parent / "v1" / "cli" / "templates" / "world-runner.Dockerfile"
@@ -253,8 +179,6 @@ def _get_docker_platform(override: str | None = None) -> str:
     if override:
         return override
 
-    import platform as plat
-
     system = plat.system()
     machine = plat.machine().lower()
 
@@ -424,8 +348,6 @@ async def _create_chronos_session(
     tags: list[str] | None = None,
 ) -> dict:
     """Create a session in Chronos."""
-    import httpx
-
     url = f"{chronos_url.rstrip('/')}/api/sessions"
 
     async with httpx.AsyncClient(timeout=30.0) as client:
@@ -443,25 +365,6 @@ async def _create_chronos_session(
         return response.json()
 
 
-async def _close_chronos_session(
-    chronos_url: str,
-    api_key: str,
-    session_id: str,
-) -> None:
-    """Close a Chronos session."""
-    import httpx
-
-    url = f"{chronos_url.rstrip('/')}/api/sessions/{session_id}/close"
-
-    try:
-        async with httpx.AsyncClient(timeout=30.0) as client:
-            response = await client.post(url, headers={"x-api-key": api_key})
-            response.raise_for_status()
-            logger.info(f"Closed Chronos session: {session_id}")
-    except Exception as e:
-        logger.warning(f"Failed to close Chronos session: {e}")
-
-
 async def _complete_chronos_session(
     chronos_url: str,
     api_key: str,
@@ -471,7 +374,6 @@ async def _complete_chronos_session(
     error_message: str | None = None,
 ) -> None:
     """Complete a Chronos session with final status."""
-    import httpx
 
     url = f"{chronos_url.rstrip('/')}/api/sessions/{session_id}/complete"
 

plato/worlds/__init__.py

@@ -46,23 +46,19 @@ from plato._generated.models import (
 from plato.worlds.base import (
     BaseWorld,
     ConfigT,
-    Observation,
-    StepResult,
     get_registered_worlds,
     get_world,
     register_world,
 )
 from plato.worlds.config import (
-    Agent,
     AgentConfig,
     CheckpointConfig,
-    Env,
     EnvConfig,
-    EnvList,
     RunConfig,
-    Secret,
     StateConfig,
 )
+from plato.worlds.markers import Agent, Env, EnvList, Secret
+from plato.worlds.models import Observation, StepResult
 from plato.worlds.runner import run_world
 
 __all__ = [