plato-sdk-v2 2.3.0__py3-none-any.whl → 2.4.1__py3-none-any.whl

plato/agents/__init__.py

@@ -20,8 +20,9 @@ Trajectory (ATIF):
     - Trajectory: ATIF trajectory model
     - Step, Agent, ToolCall, etc.: ATIF components
 
-Callback:
-    - ChronosCallback: Utility for Chronos communication
+OTel Tracing:
+    - instrument: Initialize OTel tracing from environment
+    - get_tracer: Get a tracer for creating spans
 
 Example (direct execution):
     from plato.agents import BaseAgent, AgentConfig, Secret, register_agent
@@ -80,14 +81,25 @@ __all__ = [
     "Metrics",
     "FinalMetrics",
     "SCHEMA_VERSION",
-    # Logging
-    "init_logging",
-    "span",
-    "log_event",
+    # Artifacts
+    "zip_directory",
     "upload_artifacts",
-    "reset_logging",
+    "upload_artifact",
+    "upload_to_s3",
+    # OTel tracing
+    "init_tracing",
+    "instrument",
+    "shutdown_tracing",
+    "get_tracer",
+    "is_initialized",
 ]
 
+from plato.agents.artifacts import (
+    upload_artifact,
+    upload_artifacts,
+    upload_to_s3,
+    zip_directory,
+)
 from plato.agents.base import (
     BaseAgent,
     ConfigT,
@@ -97,12 +109,12 @@ from plato.agents.base import (
 )
 from plato.agents.build import BuildConfig, load_build_config
 from plato.agents.config import AgentConfig, Secret
-from plato.agents.logging import (
-    init_logging,
-    log_event,
-    reset_logging,
-    span,
-    upload_artifacts,
+from plato.agents.otel import (
+    get_tracer,
+    init_tracing,
+    instrument,
+    is_initialized,
+    shutdown_tracing,
 )
 from plato.agents.runner import run_agent
 from plato.agents.trajectory import (

plato/agents/artifacts.py

@@ -0,0 +1,108 @@
+"""Artifact upload utilities for Plato agents and worlds.
+
+These functions upload artifacts directly to S3 using presigned URLs.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import zipfile
+from pathlib import Path
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+
+def zip_directory(dir_path: str) -> bytes:
+    """Zip an entire directory.
+
+    Args:
+        dir_path: Path to the directory
+
+    Returns:
+        Zip file contents as bytes.
+    """
+    path = Path(dir_path)
+    buffer = io.BytesIO()
+
+    with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zf:
+        for file_path in path.rglob("*"):
+            if file_path.is_file():
+                arcname = file_path.relative_to(path)
+                zf.write(file_path, arcname)
+
+    buffer.seek(0)
+    return buffer.read()
+
+
+async def upload_to_s3(upload_url: str, data: bytes, content_type: str = "application/octet-stream") -> bool:
+    """Upload data directly to S3 using a presigned URL.
+
+    Args:
+        upload_url: Presigned S3 PUT URL
+        data: Raw bytes to upload
+        content_type: MIME type of the content
+
+    Returns:
+        True if successful, False otherwise
+    """
+    if not upload_url:
+        logger.warning("No upload URL provided")
+        return False
+
+    try:
+        async with httpx.AsyncClient(timeout=120.0) as client:
+            response = await client.put(
+                upload_url,
+                content=data,
+                headers={"Content-Type": content_type},
+            )
+            if response.status_code in (200, 201, 204):
+                logger.info(f"Uploaded {len(data)} bytes to S3")
+                return True
+            else:
+                logger.warning(f"S3 upload failed: {response.status_code} {response.text}")
+                return False
+    except Exception as e:
+        logger.warning(f"Failed to upload to S3: {e}")
+        return False
+
+
+async def upload_artifacts(upload_url: str, dir_path: str) -> bool:
+    """Upload a directory as a zip directly to S3.
+
+    Args:
+        upload_url: Presigned S3 PUT URL
+        dir_path: Path to the directory to upload
+
+    Returns:
+        True if successful, False otherwise
+    """
+    try:
+        zip_data = zip_directory(dir_path)
+        logger.info(f"Zipped directory: {len(zip_data)} bytes")
+    except Exception as e:
+        logger.warning(f"Failed to zip directory: {e}")
+        return False
+
+    return await upload_to_s3(upload_url, zip_data, "application/zip")
+
+
+async def upload_artifact(
+    upload_url: str,
+    data: bytes,
+    content_type: str = "application/octet-stream",
+) -> bool:
+    """Upload an artifact directly to S3.
+
+    Args:
+        upload_url: Presigned S3 PUT URL
+        data: Raw bytes of the artifact
+        content_type: MIME type of the content
+
+    Returns:
+        True if successful, False otherwise
+    """
+    return await upload_to_s3(upload_url, data, content_type)

plato/agents/config.py

@@ -18,10 +18,8 @@ Example:
 from __future__ import annotations
 
 import json
-from pathlib import Path
 from typing import Any
 
-from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
@@ -57,8 +55,6 @@ class AgentConfig(BaseSettings):
 
     Attributes:
         logs_dir: Directory for agent logs and trajectory output.
-        checkpoint_paths: Directories to watch for checkpoint triggers (for workspace tracking).
-        checkpoint_debounce_ms: Debounce interval for checkpoints.
     """
 
     model_config = SettingsConfigDict(
@@ -68,8 +64,6 @@ class AgentConfig(BaseSettings):
     )
 
     logs_dir: str = "/logs"
-    checkpoint_paths: list[str] = Field(default_factory=list)
-    checkpoint_debounce_ms: int = 500
 
     @classmethod
     def get_field_secrets(cls) -> dict[str, Secret]:
@@ -97,7 +91,7 @@ class AgentConfig(BaseSettings):
         secrets = []
 
         # Skip internal fields
-        internal_fields = {"logs_dir", "checkpoint_paths", "checkpoint_debounce_ms"}
+        internal_fields = {"logs_dir"}
 
         for field_name, prop_schema in properties.items():
             if field_name in internal_fields:
@@ -140,7 +134,7 @@ class AgentConfig(BaseSettings):
     def get_config_dict(self) -> dict[str, Any]:
         """Extract non-secret config values as a dict."""
         secrets_map = self.get_field_secrets()
-        internal_fields = {"logs_dir", "checkpoint_paths", "checkpoint_debounce_ms"}
+        internal_fields = {"logs_dir"}
 
         result: dict[str, Any] = {}
         for field_name in self.model_fields:
@@ -152,9 +146,18 @@ class AgentConfig(BaseSettings):
         return result
 
     @classmethod
-    def from_file(cls, path: str | Path) -> AgentConfig:
-        """Load config from a JSON file."""
-        path = Path(path)
-        with open(path) as f:
-            data = json.load(f)
+    def from_env(cls) -> AgentConfig:
+        """Load config from AGENT_CONFIG_B64 environment variable.
+
+        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")
+        config_json = base64.b64decode(config_b64).decode()
+        data = json.loads(config_json)
         return cls(**data)

plato/agents/otel.py

@@ -0,0 +1,261 @@
+"""OpenTelemetry integration for Plato agents and worlds.
+
+Provides tracing utilities using OpenTelemetry SDK. Traces are sent directly
+to the Chronos OTLP endpoint.
+
+Usage:
+    from plato.agents.otel import init_tracing, get_tracer, shutdown_tracing
+
+    # Initialize tracing (sends to Chronos OTLP endpoint)
+    init_tracing(
+        service_name="my-world",
+        session_id="session-123",
+        otlp_endpoint="http://chronos/api/otel",
+    )
+
+    # Create spans
+    tracer = get_tracer()
+    with tracer.start_as_current_span("my-operation") as span:
+        span.set_attribute("key", "value")
+        # ... do work ...
+
+    # Cleanup
+    shutdown_tracing()
+"""
+
+from __future__ import annotations
+
+import logging
+
+from opentelemetry import trace
+from opentelemetry.trace import Tracer
+
+_module_logger = logging.getLogger(__name__)
+
+# Global state
+_tracer_provider = None
+_initialized = False
+_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.
+    """
+
+    def __init__(self, tracer: Tracer, level: int = logging.INFO):
+        super().__init__(level)
+        self.tracer = tracer
+
+    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())
+                span.set_attribute("log.logger", record.name)
+                span.set_attribute("source", "world")
+                span.set_attribute("content", record.getMessage()[:1000])
+
+                if record.funcName:
+                    span.set_attribute("log.function", record.funcName)
+                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
+
+
+def init_tracing(
+    service_name: str,
+    session_id: str,
+    otlp_endpoint: str,
+    parent_trace_id: str | None = None,
+    parent_span_id: str | None = None,
+) -> None:
+    """Initialize OpenTelemetry tracing.
+
+    Args:
+        service_name: Name of the service (e.g., world name or agent name)
+        session_id: Chronos session ID (added as resource attribute)
+        otlp_endpoint: Chronos OTLP endpoint (e.g., http://chronos/api/otel)
+        parent_trace_id: Optional parent trace ID for linking (hex string)
+        parent_span_id: Optional parent span ID for linking (hex string)
+    """
+    global _tracer_provider, _initialized, _log_handler
+
+    if _initialized:
+        _module_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,
+                "plato.session.id": session_id,
+            }
+        )
+
+        # 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
+            )
+            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}")
+
+        # 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}")
+
+    except ImportError as e:
+        print(f"[OTel] OpenTelemetry SDK not installed: {e}")
+        _module_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}")
+
+
+def shutdown_tracing(timeout_millis: int = 30000) -> None:
+    """Shutdown the tracer provider and flush spans.
+
+    Args:
+        timeout_millis: Timeout in milliseconds to wait for flush (default 30s)
+    """
+    global _tracer_provider, _initialized, _log_handler
+
+    # Remove log handler
+    if _log_handler:
+        try:
+            plato_logger = logging.getLogger("plato")
+            plato_logger.removeHandler(_log_handler)
+        except Exception:
+            pass
+        _log_handler = 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")
+
+            _tracer_provider.shutdown()
+            print("[OTel] Tracing shutdown complete")
+        except Exception as e:
+            print(f"[OTel] Error shutting down tracer: {e}")
+            _module_logger.warning(f"Error shutting down tracer: {e}")
+
+    _tracer_provider = None
+    _initialized = False
+
+
+def get_tracer(name: str = "plato") -> Tracer:
+    """Get a tracer instance.
+
+    Args:
+        name: Tracer name (default: "plato")
+
+    Returns:
+        OpenTelemetry Tracer
+    """
+    return trace.get_tracer(name)
+
+
+def is_initialized() -> bool:
+    """Check if OTel tracing is initialized."""
+    return _initialized
+
+
+def instrument(service_name: str = "plato-agent") -> Tracer:
+    """Initialize OTel tracing from environment variables.
+
+    Reads the following env vars:
+    - OTEL_EXPORTER_OTLP_ENDPOINT: Chronos OTLP endpoint (required for tracing)
+    - SESSION_ID: Chronos session ID (default: "local")
+    - OTEL_TRACE_ID: Parent trace ID for linking spans (optional)
+    - OTEL_PARENT_SPAN_ID: Parent span ID for linking spans (optional)
+
+    If OTEL_EXPORTER_OTLP_ENDPOINT is not set, returns a no-op tracer.
+
+    Args:
+        service_name: Name of the service for traces
+
+    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")
+        return trace.get_tracer(service_name)
+
+    # Initialize tracing with parent context if provided
+    init_tracing(
+        service_name=service_name,
+        session_id=session_id,
+        otlp_endpoint=otel_endpoint,
+        parent_trace_id=parent_trace_id,
+        parent_span_id=parent_span_id,
+    )
+
+    return trace.get_tracer(service_name)