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

plato/agents/callback.py

@@ -1,246 +0,0 @@
-"""Chronos callback utilities for agents.
-
-This module provides utilities for agents to communicate with Chronos,
-including pushing logs, uploading trajectories, and uploading zipped logs.
-"""
-
-from __future__ import annotations
-
-import base64
-import io
-import json
-import logging
-import zipfile
-from pathlib import Path
-from typing import Any
-
-import httpx
-
-logger = logging.getLogger(__name__)
-
-
-class ChronosCallback:
-    """Utility class for communicating with Chronos server.
-
-    Handles pushing logs, updating status, and uploading artifacts
-    (trajectory and zipped logs) to the Chronos callback endpoints.
-
-    Example:
-        callback = ChronosCallback(
-            callback_url="http://chronos.example.com/api/callback",
-            session_id="abc123",
-        )
-
-        # Push logs during execution
-        await callback.push_logs([
-            {"level": "info", "message": "Starting agent..."},
-        ])
-
-        # After agent completes, upload artifacts
-        await callback.upload_artifacts(logs_dir="/path/to/logs")
-    """
-
-    def __init__(self, callback_url: str, session_id: str):
-        """Initialize the callback client.
-
-        Args:
-            callback_url: Full callback base URL (e.g., http://server/api/callback)
-            session_id: The Chronos session ID for this run
-        """
-        self.callback_url = callback_url.rstrip("/")
-        self.session_id = session_id
-        self._enabled = bool(callback_url and session_id)
-
-    @property
-    def enabled(self) -> bool:
-        """Check if callbacks are enabled (both server and session_id set)."""
-        return self._enabled
-
-    async def push_logs(self, logs: list[dict[str, Any]]) -> bool:
-        """Push log entries to Chronos.
-
-        Args:
-            logs: List of log entries, each with 'level' and 'message' keys.
-                  Optional keys: 'timestamp', 'extra'
-
-        Returns:
-            True if logs were pushed successfully, False otherwise.
-        """
-        if not self._enabled:
-            return False
-
-        try:
-            async with httpx.AsyncClient(timeout=10.0) as client:
-                response = await client.post(
-                    f"{self.callback_url}/logs",
-                    json={"session_id": self.session_id, "logs": logs},
-                )
-                return response.status_code == 200
-        except Exception as e:
-            logger.warning(f"Failed to push logs to Chronos: {e}")
-            return False
-
-    async def push_log(
-        self,
-        message: str,
-        level: str = "info",
-        extra: dict[str, Any] | None = None,
-    ) -> bool:
-        """Push a single log entry to Chronos.
-
-        Args:
-            message: Log message
-            level: Log level (debug, info, warning, error)
-            extra: Optional additional structured data
-
-        Returns:
-            True if log was pushed successfully, False otherwise.
-        """
-        log_entry: dict[str, Any] = {"level": level, "message": message}
-        if extra:
-            log_entry["extra"] = extra
-        return await self.push_logs([log_entry])
-
-    async def update_status(
-        self,
-        status: str,
-        message: str | None = None,
-        extra: dict[str, Any] | None = None,
-    ) -> bool:
-        """Update the session status in Chronos.
-
-        Args:
-            status: New status (running, completed, failed)
-            message: Optional status message
-            extra: Optional additional data
-
-        Returns:
-            True if status was updated successfully, False otherwise.
-        """
-        if not self._enabled:
-            return False
-
-        try:
-            async with httpx.AsyncClient(timeout=10.0) as client:
-                response = await client.post(
-                    f"{self.callback_url}/status",
-                    json={
-                        "session_id": self.session_id,
-                        "status": status,
-                        "message": message,
-                        "extra": extra,
-                    },
-                )
-                return response.status_code == 200
-        except Exception as e:
-            logger.warning(f"Failed to update status in Chronos: {e}")
-            return False
-
-    def find_trajectory(self, logs_dir: str) -> dict[str, Any] | None:
-        """Find and load ATIF trajectory from logs directory.
-
-        Agents write ATIF trajectory to /logs/agent/trajectory.json after
-        completing their run. This method finds and loads that file.
-
-        Args:
-            logs_dir: Path to the logs directory
-
-        Returns:
-            Parsed ATIF trajectory dict if found, None otherwise.
-        """
-        trajectory_path = Path(logs_dir) / "agent" / "trajectory.json"
-
-        if not trajectory_path.exists():
-            return None
-
-        try:
-            with open(trajectory_path) as f:
-                data = json.load(f)
-                if isinstance(data, dict) and "schema_version" in data:
-                    logger.info(f"Found ATIF trajectory: {trajectory_path}")
-                    return data
-                return None
-        except Exception as e:
-            logger.warning(f"Failed to load trajectory from {trajectory_path}: {e}")
-            return None
-
-    def zip_logs_dir(self, logs_dir: str) -> bytes:
-        """Zip the entire logs directory.
-
-        Args:
-            logs_dir: Path to the logs directory
-
-        Returns:
-            Zip file contents as bytes.
-        """
-        logs_path = Path(logs_dir)
-        buffer = io.BytesIO()
-
-        with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zf:
-            for file_path in logs_path.rglob("*"):
-                if file_path.is_file():
-                    arcname = file_path.relative_to(logs_path)
-                    zf.write(file_path, arcname)
-
-        buffer.seek(0)
-        return buffer.read()
-
-    async def upload_artifacts(
-        self,
-        logs_dir: str,
-        trajectory: dict[str, Any] | None = None,
-    ) -> bool:
-        """Upload trajectory and zipped logs to Chronos.
-
-        If trajectory is not provided, attempts to find it in the logs directory.
-        The logs directory is zipped and uploaded to S3 via Chronos.
-
-        Args:
-            logs_dir: Path to the logs directory
-            trajectory: Optional pre-loaded trajectory dict. If None, will
-                       attempt to find trajectory.json in logs_dir.
-
-        Returns:
-            True if artifacts were uploaded successfully, False otherwise.
-        """
-        if not self._enabled:
-            return False
-
-        # Find trajectory if not provided
-        if trajectory is None:
-            trajectory = self.find_trajectory(logs_dir)
-            if trajectory:
-                logger.info("Found ATIF trajectory in logs directory")
-            else:
-                logger.info("No ATIF trajectory found in logs directory")
-
-        # Zip logs directory
-        logs_base64: str | None = None
-        try:
-            logs_zip = self.zip_logs_dir(logs_dir)
-            logs_base64 = base64.b64encode(logs_zip).decode("utf-8")
-            logger.info(f"Zipped logs: {len(logs_zip)} bytes")
-        except Exception as e:
-            logger.warning(f"Failed to zip logs: {e}")
-
-        # Upload to Chronos
-        try:
-            async with httpx.AsyncClient(timeout=60.0) as client:
-                response = await client.post(
-                    f"{self.callback_url}/artifacts",
-                    json={
-                        "session_id": self.session_id,
-                        "trajectory": trajectory,
-                        "logs_base64": logs_base64,
-                    },
-                )
-                if response.status_code == 200:
-                    result = response.json()
-                    logger.info(f"Uploaded artifacts to Chronos: {result}")
-                    return True
-                else:
-                    logger.warning(f"Failed to upload artifacts: {response.status_code} {response.text}")
-                    return False
-        except Exception as e:
-            logger.warning(f"Failed to upload artifacts to Chronos: {e}")
-            return False