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

plato/agents/logging.py

@@ -1,515 +0,0 @@
-"""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
-
-
-# =============================================================================
-# Artifact Upload (Generic)
-# =============================================================================
-
-
-async def upload_artifact(
-    data: bytes,
-    artifact_type: str,
-    filename: str | None = None,
-    extra: dict[str, Any] | None = None,
-) -> dict[str, Any] | None:
-    """Upload an artifact to Chronos.
-
-    Artifacts are stored in S3 and linked to the session in the database.
-
-    Args:
-        data: Raw bytes of the artifact
-        artifact_type: Type of artifact (e.g., "state", "logs", "trajectory")
-        filename: Optional filename for the artifact
-        extra: Optional extra data to store with the artifact
-
-    Returns:
-        Dict with artifact_id and s3_url if successful, None otherwise.
-    """
-    chronos = _ChronosLogger._instance
-    if not chronos or not chronos.enabled:
-        return None
-
-    try:
-        data_base64 = base64.b64encode(data).decode("utf-8")
-        logger.info(f"Uploading artifact: type={artifact_type}, size={len(data)} bytes")
-    except Exception as e:
-        logger.warning(f"Failed to encode artifact: {e}")
-        return None
-
-    try:
-        async with httpx.AsyncClient(timeout=120.0) as client:
-            response = await client.post(
-                f"{chronos.callback_url}/artifact",
-                json={
-                    "session_id": chronos.session_id,
-                    "artifact_type": artifact_type,
-                    "data_base64": data_base64,
-                    "filename": filename,
-                    "extra": extra or {},
-                },
-            )
-            if response.status_code == 200:
-                result = response.json()
-                logger.info(f"Uploaded artifact: {result}")
-                return result
-            else:
-                logger.warning(f"Failed to upload artifact: {response.status_code} {response.text}")
-                return None
-    except Exception as e:
-        logger.warning(f"Failed to upload artifact: {e}")
-        return None
-
-
-# =============================================================================
-# Checkpoint Upload
-# =============================================================================
-
-
-async def upload_checkpoint(
-    step_number: int,
-    env_snapshots: dict[str, str],
-    state_artifact_id: str | None = None,
-    extra: dict[str, Any] | None = None,
-) -> dict[str, Any] | None:
-    """Upload checkpoint data to Chronos.
-
-    A checkpoint includes:
-    - Environment snapshots (artifact IDs per env alias)
-    - State artifact (git bundle of /state directory)
-    - Extra data (step number, timestamp, etc.)
-
-    Args:
-        step_number: The step number when this checkpoint was created
-        env_snapshots: Dict mapping env alias to artifact_id
-        state_artifact_id: Artifact ID of the state bundle (from upload_artifact)
-        extra: Optional additional data
-
-    Returns:
-        Dict with checkpoint_id if successful, None otherwise.
-    """
-    chronos = _ChronosLogger._instance
-    if not chronos or not chronos.enabled:
-        return None
-
-    try:
-        async with httpx.AsyncClient(timeout=60.0) as client:
-            response = await client.post(
-                f"{chronos.callback_url}/checkpoint",
-                json={
-                    "session_id": chronos.session_id,
-                    "step_number": step_number,
-                    "env_snapshots": env_snapshots,
-                    "state_artifact_id": state_artifact_id,
-                    "extra": extra or {},
-                },
-            )
-            if response.status_code == 200:
-                result = response.json()
-                logger.info(f"Uploaded checkpoint: step={step_number}, checkpoint_id={result.get('checkpoint_id')}")
-                return result
-            else:
-                logger.warning(f"Failed to upload checkpoint: {response.status_code} {response.text}")
-                return None
-    except Exception as e:
-        logger.warning(f"Failed to upload checkpoint: {e}")
-        return None

plato/chronos/api/callback/__init__.py

@@ -1,11 +0,0 @@
-"""API endpoints."""
-
-from . import push_agent_logs, update_agent_status, upload_artifacts, upload_logs_zip, upload_trajectory
-
-__all__ = [
-    "push_agent_logs",
-    "update_agent_status",
-    "upload_trajectory",
-    "upload_logs_zip",
-    "upload_artifacts",
-]

plato/chronos/api/callback/push_agent_logs.py

@@ -1,61 +0,0 @@
-"""Push Agent Logs"""
-
-from __future__ import annotations
-
-from typing import Any
-
-import httpx
-
-from plato.chronos.errors import raise_for_status
-from plato.chronos.models import AgentLogsRequest, AgentLogsResponse
-
-
-def _build_request_args(
-    body: AgentLogsRequest,
-) -> dict[str, Any]:
-    """Build request arguments."""
-    url = "/api/callback/logs"
-
-    return {
-        "method": "POST",
-        "url": url,
-        "json": body.model_dump(mode="json", exclude_none=True),
-    }
-
-
-def sync(
-    client: httpx.Client,
-    body: AgentLogsRequest,
-) -> AgentLogsResponse:
-    """Receive logs from an agent running in a VM.
-
-    This endpoint acknowledges log receipt for real-time streaming.
-    Logs are NOT stored in the database - they're uploaded to S3 as a zip
-    at the end of the session via the /artifacts endpoint."""
-
-    request_args = _build_request_args(
-        body=body,
-    )
-
-    response = client.request(**request_args)
-    raise_for_status(response)
-    return AgentLogsResponse.model_validate(response.json())
-
-
-async def asyncio(
-    client: httpx.AsyncClient,
-    body: AgentLogsRequest,
-) -> AgentLogsResponse:
-    """Receive logs from an agent running in a VM.
-
-    This endpoint acknowledges log receipt for real-time streaming.
-    Logs are NOT stored in the database - they're uploaded to S3 as a zip
-    at the end of the session via the /artifacts endpoint."""
-
-    request_args = _build_request_args(
-        body=body,
-    )
-
-    response = await client.request(**request_args)
-    raise_for_status(response)
-    return AgentLogsResponse.model_validate(response.json())

plato/chronos/api/callback/update_agent_status.py

@@ -1,57 +0,0 @@
-"""Update Agent Status"""
-
-from __future__ import annotations
-
-from typing import Any
-
-import httpx
-
-from plato.chronos.errors import raise_for_status
-from plato.chronos.models import AgentLogsResponse, AgentStatusRequest
-
-
-def _build_request_args(
-    body: AgentStatusRequest,
-) -> dict[str, Any]:
-    """Build request arguments."""
-    url = "/api/callback/status"
-
-    return {
-        "method": "POST",
-        "url": url,
-        "json": body.model_dump(mode="json", exclude_none=True),
-    }
-
-
-def sync(
-    client: httpx.Client,
-    body: AgentStatusRequest,
-) -> AgentLogsResponse:
-    """Update the status of a running session.
-
-    Called by agents to report completion or failure."""
-
-    request_args = _build_request_args(
-        body=body,
-    )
-
-    response = client.request(**request_args)
-    raise_for_status(response)
-    return AgentLogsResponse.model_validate(response.json())
-
-
-async def asyncio(
-    client: httpx.AsyncClient,
-    body: AgentStatusRequest,
-) -> AgentLogsResponse:
-    """Update the status of a running session.
-
-    Called by agents to report completion or failure."""
-
-    request_args = _build_request_args(
-        body=body,
-    )
-
-    response = await client.request(**request_args)
-    raise_for_status(response)
-    return AgentLogsResponse.model_validate(response.json())

plato/chronos/api/callback/upload_artifacts.py

@@ -1,59 +0,0 @@
-"""Upload Artifacts Endpoint"""
-
-from __future__ import annotations
-
-from typing import Any
-
-import httpx
-
-from plato.chronos.errors import raise_for_status
-from plato.chronos.models import ArtifactsUploadRequest, ArtifactsUploadResponse
-
-
-def _build_request_args(
-    body: ArtifactsUploadRequest,
-) -> dict[str, Any]:
-    """Build request arguments."""
-    url = "/api/callback/artifacts"
-
-    return {
-        "method": "POST",
-        "url": url,
-        "json": body.model_dump(mode="json", exclude_none=True),
-    }
-
-
-def sync(
-    client: httpx.Client,
-    body: ArtifactsUploadRequest,
-) -> ArtifactsUploadResponse:
-    """Upload trajectory and/or logs for a session.
-
-    Convenience endpoint to upload both artifacts in one request.
-    Trajectory is stored in DB, logs are uploaded to S3."""
-
-    request_args = _build_request_args(
-        body=body,
-    )
-
-    response = client.request(**request_args)
-    raise_for_status(response)
-    return ArtifactsUploadResponse.model_validate(response.json())
-
-
-async def asyncio(
-    client: httpx.AsyncClient,
-    body: ArtifactsUploadRequest,
-) -> ArtifactsUploadResponse:
-    """Upload trajectory and/or logs for a session.
-
-    Convenience endpoint to upload both artifacts in one request.
-    Trajectory is stored in DB, logs are uploaded to S3."""
-
-    request_args = _build_request_args(
-        body=body,
-    )
-
-    response = await client.request(**request_args)
-    raise_for_status(response)
-    return ArtifactsUploadResponse.model_validate(response.json())