plato-sdk-v2 2.1.14__tar.gz → 2.1.16__tar.gz

{plato_sdk_v2-2.1.14 → plato_sdk_v2-2.1.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plato-sdk-v2
-Version: 2.1.14
+Version: 2.1.16
 Summary: Python SDK for the Plato API
 Author-email: Plato <support@plato.so>
 License-Expression: MIT

{plato_sdk_v2-2.1.14 → plato_sdk_v2-2.1.16}/plato/agents/__init__.py

@@ -69,8 +69,7 @@ __all__ = [
     "get_agent",
     "get_registered_agents",
     # Runner
-    "AgentRunner",
-    "AgentRunResult",
+    "run_agent",
     # Trajectory (ATIF)
     "Trajectory",
     "Step",
@@ -81,8 +80,12 @@ __all__ = [
     "Metrics",
     "FinalMetrics",
     "SCHEMA_VERSION",
-    # Callback
-    "ChronosCallback",
+    # Logging
+    "init_logging",
+    "span",
+    "log_event",
+    "upload_artifacts",
+    "reset_logging",
 ]
 
 from plato.agents.base import (
@@ -93,9 +96,15 @@ from plato.agents.base import (
     register_agent,
 )
 from plato.agents.build import BuildConfig, load_build_config
-from plato.agents.callback import ChronosCallback
 from plato.agents.config import AgentConfig, Secret
-from plato.agents.runner import AgentRunner, AgentRunResult
+from plato.agents.logging import (
+    init_logging,
+    log_event,
+    reset_logging,
+    span,
+    upload_artifacts,
+)
+from plato.agents.runner import run_agent
 from plato.agents.trajectory import (
     SCHEMA_VERSION,
     Agent,

plato_sdk_v2-2.1.16/plato/agents/logging.py

@@ -0,0 +1,401 @@
+"""Chronos logging utilities.
+
+Simple singleton-based logging for Chronos events and spans.
+
+Usage:
+    from plato.agents.logging import init_logging, span, log_event, upload_logs
+
+    # Initialize once (typically in world setup)
+    init_logging(
+        callback_url="http://chronos/api/callback",
+        session_id="session-123",
+    )
+
+    # Use spans anywhere - parent tracking is automatic
+    async with span("clone_repo") as s:
+        s.log("Cloning...")
+
+        async with span("checkout"):  # Automatically nested under clone_repo
+            s.log("Checking out...")
+
+    # Log events - parent is automatically the current span
+    await log_event(span_type="my_event", content="Something happened")
+
+    # Upload artifacts (zip of directory)
+    await upload_artifacts(dir_path="/path/to/logs")
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import logging
+import zipfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+
+class _ChronosLogger:
+    """Internal singleton logger for Chronos events."""
+
+    _instance: _ChronosLogger | None = None
+
+    def __init__(
+        self,
+        callback_url: str,
+        session_id: str,
+        parent_event_id: str | None = None,
+    ):
+        self.callback_url = callback_url.rstrip("/")
+        self.session_id = session_id
+        self.parent_event_id = parent_event_id
+        self._current_span_id: str | None = None
+        self._enabled = bool(callback_url and session_id)
+
+    @property
+    def enabled(self) -> bool:
+        return self._enabled
+
+    @property
+    def current_parent_id(self) -> str | None:
+        """Get the current parent ID (current span or root parent)."""
+        return self._current_span_id or self.parent_event_id
+
+    async def log_event(
+        self,
+        span_type: str,
+        content: str = "",
+        source: str = "agent",
+        log_type: str | None = None,
+        step_number: int | None = None,
+        extra: dict[str, Any] | None = None,
+        event_id: str | None = None,
+        parent_id: str | None = None,
+        started_at: str | None = None,
+        ended_at: str | None = None,
+    ) -> bool:
+        """Log an event to Chronos."""
+        if not self._enabled:
+            return False
+
+        try:
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.post(
+                    f"{self.callback_url}/event",
+                    json={
+                        "session_id": self.session_id,
+                        "span_type": span_type,
+                        "content": content,
+                        "source": source,
+                        "log_type": log_type,
+                        "step_number": step_number,
+                        "extra": extra,
+                        "event_id": event_id,
+                        "parent_id": parent_id,
+                        "started_at": started_at,
+                        "ended_at": ended_at,
+                    },
+                )
+                if response.status_code == 200:
+                    return True
+                else:
+                    logger.warning(f"Failed to log event: {response.status_code} {response.text}")
+                    return False
+        except Exception as e:
+            logger.warning(f"Failed to log event to Chronos: {e}")
+            return False
+
+
+class SpanContext:
+    """Context manager for logging spans with automatic timing and nesting."""
+
+    def __init__(
+        self,
+        name: str,
+        span_type: str = "span",
+        source: str = "agent",
+        extra: dict[str, Any] | None = None,
+    ):
+        self._name = name
+        self._span_type = span_type
+        self._source = source
+        self._extra = extra or {}
+        self._event_id: str | None = None
+        self._started_at: datetime | None = None
+        self._child_logs: list[str] = []
+        self._previous_span_id: str | None = None
+
+    async def __aenter__(self) -> SpanContext:
+        """Log span start and set as current span."""
+        self._started_at = datetime.now(timezone.utc)
+        self._event_id = str(uuid4())
+
+        chronos = _ChronosLogger._instance
+        if chronos and chronos.enabled:
+            # Save previous span and set self as current
+            self._previous_span_id = chronos._current_span_id
+            parent_id = chronos.current_parent_id
+
+            await chronos.log_event(
+                span_type=self._span_type,
+                content=f"{self._name} started",
+                source=self._source,
+                event_id=self._event_id,
+                parent_id=parent_id,
+                started_at=self._started_at.isoformat(),
+                extra=self._extra,
+            )
+
+            # Set self as the current span for child events
+            chronos._current_span_id = self._event_id
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Log span end and restore previous span."""
+        ended_at = datetime.now(timezone.utc)
+        status = "failed" if exc_type else "completed"
+
+        chronos = _ChronosLogger._instance
+        if chronos and chronos.enabled:
+            # Restore previous span
+            chronos._current_span_id = self._previous_span_id
+            parent_id = self._previous_span_id or chronos.parent_event_id
+
+            # Update extra with any child logs
+            extra = {**self._extra}
+            if self._child_logs:
+                extra["logs"] = self._child_logs
+            if exc_type:
+                extra["error"] = str(exc_val)
+
+            await chronos.log_event(
+                span_type=self._span_type,
+                content=f"{self._name} {status}",
+                source=self._source,
+                event_id=self._event_id,
+                parent_id=parent_id,
+                started_at=self._started_at.isoformat() if self._started_at else None,
+                ended_at=ended_at.isoformat(),
+                extra=extra,
+            )
+
+    def log(self, message: str) -> None:
+        """Add a log message to the span."""
+        self._child_logs.append(message)
+        logger.info(f"[{self._name}] {message}")
+
+    def set_extra(self, data: dict[str, Any]) -> None:
+        """Update the span's extra data."""
+        self._extra.update(data)
+
+    @property
+    def event_id(self) -> str | None:
+        """Get the span's event ID (available after entering context)."""
+        return self._event_id
+
+
+class _NoOpSpanContext:
+    """No-op span context when logging is not initialized."""
+
+    def __init__(self, name: str, **kwargs):
+        self._name = name
+        self._event_id = str(uuid4())
+
+    async def __aenter__(self) -> _NoOpSpanContext:
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        pass
+
+    def log(self, message: str) -> None:
+        logger.debug(f"[{self._name}] {message}")
+
+    def set_extra(self, data: dict[str, Any]) -> None:
+        pass
+
+    @property
+    def event_id(self) -> str | None:
+        return self._event_id
+
+
+# =============================================================================
+# Public API
+# =============================================================================
+
+
+def init_logging(
+    callback_url: str,
+    session_id: str,
+    parent_event_id: str | None = None,
+) -> None:
+    """Initialize Chronos logging.
+
+    Call this once at startup (e.g., in world.run()).
+
+    Args:
+        callback_url: Chronos callback URL (e.g., http://chronos/api/callback)
+        session_id: Chronos session ID
+        parent_event_id: Optional root parent event ID (e.g., step ID)
+    """
+    _ChronosLogger._instance = _ChronosLogger(
+        callback_url=callback_url,
+        session_id=session_id,
+        parent_event_id=parent_event_id,
+    )
+    logger.info(f"Chronos logging initialized: session={session_id}")
+
+
+def set_parent_event_id(parent_event_id: str | None) -> None:
+    """Set the root parent event ID.
+
+    Useful for updating the parent when entering a new step.
+    """
+    if _ChronosLogger._instance:
+        _ChronosLogger._instance.parent_event_id = parent_event_id
+
+
+def span(
+    name: str,
+    span_type: str = "span",
+    source: str = "agent",
+    extra: dict[str, Any] | None = None,
+) -> SpanContext | _NoOpSpanContext:
+    """Create a span context manager.
+
+    Usage:
+        async with span("my_operation") as s:
+            s.log("Starting...")
+            # nested spans work automatically
+            async with span("sub_operation"):
+                ...
+
+    Args:
+        name: Span name/description
+        span_type: Event type for the span
+        source: Event source (agent, world, system)
+        extra: Additional data
+
+    Returns:
+        SpanContext (or NoOpSpanContext if logging not initialized)
+    """
+    if _ChronosLogger._instance and _ChronosLogger._instance.enabled:
+        return SpanContext(
+            name=name,
+            span_type=span_type,
+            source=source,
+            extra=extra,
+        )
+    else:
+        return _NoOpSpanContext(name=name)
+
+
+async def log_event(
+    span_type: str,
+    content: str = "",
+    source: str = "agent",
+    log_type: str | None = None,
+    extra: dict[str, Any] | None = None,
+    parent_id: str | None = None,
+) -> bool:
+    """Log a single event to Chronos.
+
+    For spans with timing, use span() instead.
+    """
+    chronos = _ChronosLogger._instance
+    if not chronos or not chronos.enabled:
+        return False
+
+    return await chronos.log_event(
+        span_type=span_type,
+        content=content,
+        source=source,
+        log_type=log_type,
+        extra=extra,
+        parent_id=parent_id or chronos.current_parent_id,
+    )
+
+
+def is_logging_enabled() -> bool:
+    """Check if Chronos logging is enabled."""
+    return _ChronosLogger._instance is not None and _ChronosLogger._instance.enabled
+
+
+def reset_logging() -> None:
+    """Reset the logger (mainly for testing)."""
+    _ChronosLogger._instance = None
+
+
+# =============================================================================
+# Logs Upload
+# =============================================================================
+
+
+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_artifacts(dir_path: str) -> str | None:
+    """Upload a directory as a zip to Chronos.
+
+    Args:
+        dir_path: Path to the directory to upload
+
+    Returns:
+        URL if successful, None otherwise.
+    """
+    chronos = _ChronosLogger._instance
+    if not chronos or not chronos.enabled:
+        return None
+
+    try:
+        zip_data = zip_directory(dir_path)
+        zip_base64 = base64.b64encode(zip_data).decode("utf-8")
+        logger.info(f"Zipped directory: {len(zip_data)} bytes")
+    except Exception as e:
+        logger.warning(f"Failed to zip directory: {e}")
+        return None
+
+    try:
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            response = await client.post(
+                f"{chronos.callback_url}/logs-upload",
+                json={
+                    "session_id": chronos.session_id,
+                    "logs_base64": zip_base64,
+                },
+            )
+            if response.status_code == 200:
+                result = response.json()
+                logger.info(f"Uploaded artifacts: {result}")
+                return result.get("logs_url")
+            else:
+                logger.warning(f"Failed to upload artifacts: {response.status_code} {response.text}")
+                return None
+    except Exception as e:
+        logger.warning(f"Failed to upload artifacts: {e}")
+        return None

plato_sdk_v2-2.1.16/plato/agents/runner.py

@@ -0,0 +1,150 @@
+"""Agent runner - run agents in Docker containers."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import tempfile
+from pathlib import Path
+
+from plato.agents.logging import log_event, span, upload_artifacts
+
+logger = logging.getLogger(__name__)
+
+
+async def run_agent(
+    image: str,
+    config: dict,
+    secrets: dict[str, str],
+    instruction: str,
+    workspace: str,
+    logs_dir: str | None = None,
+    pull: bool = True,
+) -> None:
+    """Run an agent in a Docker container.
+
+    Args:
+        image: Docker image URI
+        config: Agent configuration dict
+        secrets: Secret values (API keys, etc.)
+        instruction: Task instruction for the agent
+        workspace: Host directory to mount as /workspace
+        logs_dir: Host directory for logs (temp dir if None)
+        pull: Whether to pull the image first
+    """
+    logs_dir = logs_dir or tempfile.mkdtemp(prefix="agent_logs_")
+    agent_name = image.split("/")[-1].split(":")[0]
+
+    async with span(agent_name, span_type="agent", source="agent") as agent_span:
+        agent_span.log(f"Starting agent: {agent_name} ({image})")
+
+        # Pull image if requested
+        if pull:
+            agent_span.log(f"Pulling image: {image}")
+            pull_proc = await asyncio.create_subprocess_exec(
+                "docker",
+                "pull",
+                image,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.STDOUT,
+            )
+            await pull_proc.wait()
+
+        # Setup
+        os.makedirs(os.path.join(logs_dir, "agent"), exist_ok=True)
+        config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
+        json.dump(config, config_file)
+        config_file.close()
+
+        try:
+            # Build docker command
+            docker_cmd = ["docker", "run", "--rm"]
+
+            # Check if iptables is available for network isolation
+            try:
+                proc = await asyncio.create_subprocess_exec(
+                    "iptables",
+                    "-L",
+                    "-n",
+                    stdout=asyncio.subprocess.DEVNULL,
+                    stderr=asyncio.subprocess.DEVNULL,
+                )
+                await proc.wait()
+                has_iptables = proc.returncode == 0
+            except (FileNotFoundError, PermissionError):
+                has_iptables = False
+
+            if not has_iptables:
+                docker_cmd.extend(["--network=host", "--add-host=localhost:127.0.0.1"])
+
+            docker_cmd.extend(
+                [
+                    "-v",
+                    f"{workspace}:/workspace",
+                    "-v",
+                    f"{logs_dir}:/logs",
+                    "-v",
+                    f"{config_file.name}:/config.json:ro",
+                    "-w",
+                    "/workspace",
+                ]
+            )
+
+            for key, value in secrets.items():
+                docker_cmd.extend(["-e", f"{key.upper()}={value}"])
+
+            docker_cmd.append(image)
+            docker_cmd.extend(["--instruction", instruction])
+
+            # Run container and stream output
+            process = await asyncio.create_subprocess_exec(
+                *docker_cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.STDOUT,
+            )
+
+            # Stream output line by line
+            assert process.stdout is not None
+            while True:
+                line = await process.stdout.readline()
+                if not line:
+                    break
+                logger.info(f"[agent] {line.decode().rstrip()}")
+
+            await process.wait()
+
+            if process.returncode != 0:
+                raise RuntimeError(f"Agent failed with exit code {process.returncode}")
+
+            agent_span.log("Agent completed successfully")
+
+        finally:
+            os.unlink(config_file.name)
+
+            # Load trajectory and add to span
+            trajectory_path = Path(logs_dir) / "agent" / "trajectory.json"
+            if trajectory_path.exists():
+                try:
+                    with open(trajectory_path) as f:
+                        trajectory = json.load(f)
+                    if isinstance(trajectory, dict) and "schema_version" in trajectory:
+                        # Add agent image to trajectory
+                        agent_data = trajectory.get("agent", {})
+                        extra = agent_data.get("extra") or {}
+                        extra["image"] = image
+                        agent_data["extra"] = extra
+                        trajectory["agent"] = agent_data
+
+                        # Log trajectory as separate event
+                        await log_event(
+                            span_type="trajectory",
+                            log_type="atif",
+                            extra=trajectory,
+                            source="agent",
+                        )
+                except Exception as e:
+                    logger.warning(f"Failed to load trajectory: {e}")
+
+            await upload_artifacts(logs_dir)