plato-sdk-v2 2.3.3__py3-none-any.whl → 2.3.5__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,16 +81,25 @@ __all__ = [
     "Metrics",
     "FinalMetrics",
     "SCHEMA_VERSION",
-    # Logging
-    "init_logging",
-    "span",
-    "log_event",
+    # Artifacts
+    "zip_directory",
     "upload_artifacts",
     "upload_artifact",
-    "upload_checkpoint",
-    "reset_logging",
+    "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,
@@ -99,14 +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_artifact,
-    upload_artifacts,
-    upload_checkpoint,
+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

@@ -21,7 +21,6 @@ import json
 from pathlib import Path
 from typing import Any
 
-from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
@@ -57,8 +56,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 +65,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 +92,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 +135,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:

plato/agents/otel.py

@@ -0,0 +1,258 @@
+"""OpenTelemetry integration for Plato agents and worlds.
+
+Provides tracing and logging utilities using OpenTelemetry SDK. Traces and logs
+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 ...
+
+    # All Python logging is automatically sent to Chronos
+    import logging
+    logger = logging.getLogger(__name__)
+    logger.info("This will appear in the trajectory viewer!")
+
+    # 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
+_logging_handler = None
+_initialized = False
+
+
+class OTelLoggingHandler(logging.Handler):
+    """Logging handler that emits OTel spans for log messages.
+
+    Each log message becomes a span with:
+    - span.type: "log"
+    - log.level: DEBUG/INFO/WARNING/ERROR/CRITICAL
+    - content: the log message
+    - source: the logger name
+    """
+
+    def __init__(self, tracer_name: str = "plato.logging"):
+        super().__init__()
+        self._tracer_name = tracer_name
+        # Filter out noisy loggers
+        self._ignored_loggers = {
+            "httpx",
+            "httpcore",
+            "urllib3",
+            "asyncio",
+            "opentelemetry",
+            "plato.agents.otel",  # Avoid recursion
+        }
+
+    def emit(self, record: logging.LogRecord) -> None:
+        """Emit a log record as an OTel span."""
+        # Skip ignored loggers
+        logger_name = record.name
+        for ignored in self._ignored_loggers:
+            if logger_name.startswith(ignored):
+                return
+
+        try:
+            tracer = trace.get_tracer(self._tracer_name)
+
+            # Format the message
+            try:
+                msg = self.format(record)
+            except Exception:
+                msg = record.getMessage()
+
+            # Create a span for the log message
+            with tracer.start_as_current_span(
+                f"log.{record.levelname.lower()}",
+                end_on_exit=True,
+            ) as span:
+                span.set_attribute("span.type", "log")
+                span.set_attribute("log.level", record.levelname)
+                span.set_attribute("content", msg)
+                span.set_attribute("source", logger_name)
+
+                # Add extra context if available
+                if record.funcName:
+                    span.set_attribute("log.function", record.funcName)
+                if record.pathname:
+                    span.set_attribute("log.file", record.pathname)
+                if record.lineno:
+                    span.set_attribute("log.line", record.lineno)
+
+                # If there's an exception, record it
+                if record.exc_info and record.exc_info[1]:
+                    span.record_exception(record.exc_info[1])
+
+        except Exception:
+            # Don't let logging failures break the application
+            pass
+
+
+def init_tracing(
+    service_name: str,
+    session_id: str,
+    otlp_endpoint: str,
+    capture_logging: bool = True,
+    log_level: int = logging.INFO,
+) -> None:
+    """Initialize OpenTelemetry tracing and optionally capture Python logging.
+
+    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)
+        capture_logging: If True, install handler to capture Python logs as OTel spans
+        log_level: Minimum log level to capture (default: INFO)
+    """
+    global _tracer_provider, _logging_handler, _initialized
+
+    if _initialized:
+        _module_logger.debug("Tracing already initialized")
+        return
+
+    try:
+        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 BatchSpanProcessor
+
+        # Create resource with session ID
+        resource = Resource.create(
+            {
+                "service.name": service_name,
+                "session.id": session_id,
+            }
+        )
+
+        # Create tracer provider
+        _tracer_provider = TracerProvider(resource=resource)
+
+        # Add OTLP exporter pointing to Chronos
+        otlp_exporter = OTLPSpanExporter(endpoint=f"{otlp_endpoint.rstrip('/')}/v1/traces")
+        _tracer_provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
+
+        # Set as global tracer provider
+        trace.set_tracer_provider(_tracer_provider)
+
+        _initialized = True
+
+        # Install logging handler to capture Python logs
+        if capture_logging:
+            _logging_handler = OTelLoggingHandler()
+            _logging_handler.setLevel(log_level)
+            # Add to root logger to capture all logs
+            logging.getLogger().addHandler(_logging_handler)
+
+        # Use print to ensure this shows regardless of logging config
+        print(f"[OTel] Tracing initialized: service={service_name}, session={session_id}, endpoint={otlp_endpoint}")
+        _module_logger.info(
+            f"OTel tracing initialized: service={service_name}, "
+            f"session={session_id}, endpoint={otlp_endpoint}, "
+            f"capture_logging={capture_logging}"
+        )
+
+    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() -> None:
+    """Shutdown the tracer provider, flush spans, and remove logging handler."""
+    global _tracer_provider, _logging_handler, _initialized
+
+    # Remove logging handler first
+    if _logging_handler:
+        try:
+            logging.getLogger().removeHandler(_logging_handler)
+        except Exception:
+            pass
+        _logging_handler = None
+
+    if _tracer_provider:
+        try:
+            _tracer_provider.shutdown()
+            _module_logger.info("OTel tracing shutdown complete")
+        except Exception as 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")
+
+    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")
+
+    if not otel_endpoint:
+        # Return default tracer (no-op if no provider configured)
+        return trace.get_tracer(service_name)
+
+    # Initialize tracing
+    init_tracing(
+        service_name=service_name,
+        session_id=session_id,
+        otlp_endpoint=otel_endpoint,
+        capture_logging=True,
+    )
+
+    return trace.get_tracer(service_name)