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

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