plato-sdk-v2 2.0.64__py3-none-any.whl → 2.3.4__py3-none-any.whl

plato/agents/__init__.py

@@ -1,453 +1,122 @@
-"""Plato agent runner utilities.
+"""Plato agent framework.
 
-This module provides utilities for running coding agents in Docker containers.
-Each agent image has its own entrypoint that handles execution internally.
+Provides base classes and utilities for building and running agents.
+
+Base Classes:
+    - BaseAgent: Abstract base class for agents
+    - AgentConfig: Base configuration class
+    - Secret: Annotation marker for secrets
+
+Registry:
+    - register_agent: Decorator to register an agent
+    - get_agent: Get an agent by name
+    - get_registered_agents: Get all registered agents
 
 Runner:
-    - AgentRunner: Utility for running agents in Docker containers
+    - AgentRunner: Run agents in Docker containers
     - AgentRunResult: Async iterator for agent output
 
+Trajectory (ATIF):
+    - Trajectory: ATIF trajectory model
+    - Step, Agent, ToolCall, etc.: ATIF components
+
 Callback:
-    - ChronosCallback: Utility for communicating with Chronos server
+    - ChronosCallback: Utility for Chronos communication
+
+Example (direct execution):
+    from plato.agents import BaseAgent, AgentConfig, Secret, register_agent
+    from typing import Annotated
 
-Schemas:
-    - AGENT_SCHEMAS: JSON schemas for agent configuration
-    - get_agent_schema: Get schema for a specific agent
+    class MyAgentConfig(AgentConfig):
+        model_name: str = "anthropic/claude-sonnet-4"
+        api_key: Annotated[str, Secret(description="API key")]
 
-Example:
+    @register_agent("my-agent")
+    class MyAgent(BaseAgent[MyAgentConfig]):
+        name = "my-agent"
+        description = "My custom agent"
+
+        async def run(self, instruction: str) -> None:
+            # Agent implementation
+            ...
+
+Example (Docker execution):
     from plato.agents import AgentRunner
 
-    # Run agent with automatic Chronos callbacks
     async for line in AgentRunner.run(
-        image="us-docker.pkg.dev/plato-prod/agents/openhands:latest",
+        image="my-agent:latest",
         config={"model_name": "anthropic/claude-sonnet-4"},
-        secrets={"anthropic_api_key": "sk-..."},
-        instruction="Fix the bug in main.py",
+        secrets={"api_key": "sk-..."},
+        instruction="Fix the bug",
         workspace="/path/to/repo",
-        callback_url="http://chronos.example.com/api/callback",
-        session_id="abc123",
     ):
         print(line)
-    # Logs are automatically pushed and artifacts uploaded to Chronos
 """
 
 from __future__ import annotations
 
 __all__ = [
-    # Schemas
-    "get_agent_schema",
-    "AGENT_SCHEMAS",
+    # Config
+    "AgentConfig",
+    "Secret",
+    # Build
+    "BuildConfig",
+    "load_build_config",
+    # Base
+    "BaseAgent",
+    "ConfigT",
+    "register_agent",
+    "get_agent",
+    "get_registered_agents",
     # Runner
-    "AgentRunner",
-    "AgentRunResult",
-    # Callback
-    "ChronosCallback",
+    "run_agent",
+    # Trajectory (ATIF)
+    "Trajectory",
+    "Step",
+    "Agent",
+    "ToolCall",
+    "Observation",
+    "ObservationResult",
+    "Metrics",
+    "FinalMetrics",
+    "SCHEMA_VERSION",
+    # Logging
+    "init_logging",
+    "span",
+    "log_event",
+    "upload_artifacts",
+    "upload_artifact",
+    "upload_checkpoint",
+    "reset_logging",
 ]
 
-from plato.agents.callback import ChronosCallback
-
-# JSON Schemas for agent configuration
-AGENT_SCHEMAS: dict[str, dict] = {
-    "claude-code": {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "$id": "claude-code",
-        "title": "ClaudeCodeConfig",
-        "description": "Configuration for Claude Code agent.",
-        "type": "object",
-        "properties": {
-            "model_name": {
-                "type": "string",
-                "description": "LLM model to use (e.g., 'anthropic/claude-sonnet-4')",
-            },
-            "max_thinking_tokens": {
-                "type": ["integer", "null"],
-                "default": None,
-                "description": "Maximum tokens for extended thinking mode",
-            },
-        },
-        "required": [],
-    },
-    "openhands": {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "$id": "openhands",
-        "title": "OpenHandsConfig",
-        "description": "Configuration for OpenHands agent.",
-        "type": "object",
-        "properties": {
-            "model_name": {
-                "type": "string",
-                "description": "LLM model to use (e.g., 'anthropic/claude-sonnet-4')",
-            },
-            "disable_tool_calls": {
-                "type": "boolean",
-                "default": False,
-                "description": "Whether to disable native function calling",
-            },
-            "reasoning_effort": {
-                "type": ["string", "null"],
-                "enum": ["low", "medium", "high", None],
-                "default": "medium",
-                "description": "Reasoning effort level for the model",
-            },
-        },
-        "required": [],
-    },
-    "codex": {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "$id": "codex",
-        "title": "CodexConfig",
-        "description": "Configuration for Codex CLI agent.",
-        "type": "object",
-        "properties": {
-            "model_name": {
-                "type": "string",
-                "description": "LLM model to use (e.g., 'openai/gpt-4o')",
-            },
-        },
-        "required": [],
-    },
-    "aider": {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "$id": "aider",
-        "title": "AiderConfig",
-        "description": "Configuration for Aider agent.",
-        "type": "object",
-        "properties": {
-            "model_name": {
-                "type": "string",
-                "description": "LLM model to use",
-            },
-        },
-        "required": [],
-    },
-    "gemini-cli": {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "$id": "gemini-cli",
-        "title": "GeminiCliConfig",
-        "description": "Configuration for Gemini CLI agent.",
-        "type": "object",
-        "properties": {
-            "model_name": {
-                "type": "string",
-                "description": "LLM model to use (e.g., 'google/gemini-2.5-pro')",
-            },
-        },
-        "required": [],
-    },
-    "goose": {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "$id": "goose",
-        "title": "GooseConfig",
-        "description": "Configuration for Block Goose agent.",
-        "type": "object",
-        "properties": {
-            "model_name": {
-                "type": "string",
-                "description": "LLM model to use",
-            },
-        },
-        "required": [],
-    },
-    "swe-agent": {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "$id": "swe-agent",
-        "title": "SweAgentConfig",
-        "description": "Configuration for SWE-agent.",
-        "type": "object",
-        "properties": {
-            "model_name": {
-                "type": "string",
-                "description": "LLM model to use",
-            },
-        },
-        "required": [],
-    },
-}
-
-
-def get_agent_schema(agent_name: str) -> dict | None:
-    """Get the JSON schema for an agent.
-
-    Args:
-        agent_name: The agent name (e.g., 'claude-code', 'openhands')
-
-    Returns:
-        JSON schema dict or None if agent not found
-    """
-    return AGENT_SCHEMAS.get(agent_name)
-
-
-class AgentRunResult:
-    """Result of running an agent.
-
-    This class is an async iterator that yields output lines from the agent.
-    It also provides access to the logs directory where agent logs are stored.
-
-    If callback_url and session_id are provided, logs are automatically
-    pushed to Chronos during execution, and artifacts are uploaded after
-    the agent completes.
-
-    Example:
-        result = AgentRunner.run(...)
-        async for line in result:
-            print(line)
-        print(f"Logs at: {result.logs_dir}")
-    """
-
-    def __init__(
-        self,
-        image: str,
-        config: dict,
-        secrets: dict[str, str],
-        instruction: str,
-        workspace: str,
-        logs_dir: str | None,
-        pull: bool,
-        callback_url: str,
-        session_id: str,
-    ):
-        self._image = image
-        self._config = config
-        self._secrets = secrets
-        self._instruction = instruction
-        self._workspace = workspace
-        self._pull = pull
-
-        # Chronos callback
-        self._callback = ChronosCallback(
-            callback_url=callback_url,
-            session_id=session_id,
-        )
-
-        # Create logs dir if not provided
-        if logs_dir is None:
-            import tempfile
-
-            self._logs_dir = tempfile.mkdtemp(prefix="agent_logs_")
-        else:
-            self._logs_dir = logs_dir
-
-    @property
-    def logs_dir(self) -> str:
-        """Host path where agent logs are stored."""
-        return self._logs_dir
-
-    @property
-    def callback(self) -> ChronosCallback:
-        """The Chronos callback client (for manual use if needed)."""
-        return self._callback
-
-    def __aiter__(self):
-        return self._stream()
-
-    async def _stream(self):
-        """Stream output from the agent."""
-        import asyncio
-        import json
-        import os
-        import tempfile
-
-        # Log buffer for batching
-        log_buffer: list[dict] = []
-
-        async def flush_logs():
-            """Push buffered logs to Chronos."""
-            if not log_buffer or not self._callback.enabled:
-                return
-            await self._callback.push_logs(log_buffer.copy())
-            log_buffer.clear()
-
-        # Push start log
-        agent_name = self._image.split("/")[-1].split(":")[0]
-        await self._callback.push_log(f"Starting agent: {agent_name} ({self._image})")
-
-        # Pull the image if requested
-        if self._pull:
-            pull_proc = await asyncio.create_subprocess_exec(
-                "docker",
-                "pull",
-                self._image,
-                stdout=asyncio.subprocess.PIPE,
-                stderr=asyncio.subprocess.STDOUT,
-            )
-            assert pull_proc.stdout is not None  # stdout is set via PIPE
-            while True:
-                line = await pull_proc.stdout.readline()
-                if not line:
-                    break
-                yield f"[pull] {line.decode().rstrip()}"
-            await pull_proc.wait()
-
-        # Create agent subdirectory for logs (Harbor writes to /logs/agent/)
-        agent_logs_subdir = os.path.join(self._logs_dir, "agent")
-        os.makedirs(agent_logs_subdir, exist_ok=True)
-
-        # Write config to a temp file
-        config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
-        json.dump(self._config, config_file)
-        config_file.close()
-
-        agent_failed = False
-        error_message = ""
-
-        try:
-            # Build docker command
-            docker_cmd = [
-                "docker",
-                "run",
-                "--rm",
-                "-v",
-                f"{self._workspace}:/workspace",
-                "-v",
-                f"{self._logs_dir}:/logs",
-                "-v",
-                f"{config_file.name}:/config.json:ro",
-                "-w",
-                "/workspace",
-            ]
-
-            # Add secrets as environment variables
-            for key, value in self._secrets.items():
-                docker_cmd.extend(["-e", f"{key.upper()}={value}"])
-
-            # Add the image and instruction argument
-            docker_cmd.append(self._image)
-            docker_cmd.extend(["--instruction", self._instruction])
-
-            # Run the container
-            process = await asyncio.create_subprocess_exec(
-                *docker_cmd,
-                stdout=asyncio.subprocess.PIPE,
-                stderr=asyncio.subprocess.STDOUT,
-            )
-            assert process.stdout is not None  # stdout is set via PIPE
-
-            # Stream output
-            while True:
-                line = await process.stdout.readline()
-                if not line:
-                    break
-                decoded = line.decode().rstrip()
-                yield decoded
-
-                # Buffer logs for Chronos callback
-                if self._callback.enabled:
-                    log_buffer.append({"level": "info", "message": decoded})
-                    if len(log_buffer) >= 10:
-                        await flush_logs()
-
-            await process.wait()
-
-            if process.returncode != 0:
-                agent_failed = True
-                error_message = f"Agent failed with exit code {process.returncode}"
-        except Exception as e:
-            agent_failed = True
-            error_message = str(e)
-            raise
-        finally:
-            # Clean up config file
-            os.unlink(config_file.name)
-
-            # Flush any remaining logs
-            await flush_logs()
-
-            # Push final status and upload artifacts
-            if self._callback.enabled:
-                if agent_failed:
-                    await self._callback.push_log(error_message, level="error")
-                else:
-                    await self._callback.push_log("Agent completed successfully")
-
-                # Upload trajectory and logs
-                await self._callback.upload_artifacts(logs_dir=self._logs_dir)
-
-        if agent_failed:
-            raise RuntimeError(error_message)
-
-
-class AgentRunner:
-    """Utility for running agents in Docker containers.
-
-    Each agent image has its own entrypoint that:
-    - Reads config from /config.json
-    - Reads secrets from environment variables
-    - Takes instruction via --instruction argument
-
-    When callback_url and session_id are provided, the runner automatically:
-    - Pushes logs to Chronos during execution
-    - Uploads trajectory and zipped logs after completion
-
-    Example:
-        import asyncio
-        from plato.agents import AgentRunner
-
-        async def main():
-            async for line in AgentRunner.run(
-                image="us-docker.pkg.dev/plato-prod/agents/openhands:latest",
-                config={"model_name": "anthropic/claude-sonnet-4"},
-                secrets={"anthropic_api_key": "sk-..."},
-                instruction="Fix the bug in main.py",
-                workspace="/path/to/repo",
-                callback_url="http://chronos.example.com/api/callback",
-                session_id="abc123",
-            ):
-                print(line)
-            # Logs and artifacts are automatically uploaded to Chronos
-
-        asyncio.run(main())
-    """
-
-    @staticmethod
-    def run(
-        image: str,
-        config: dict,
-        secrets: dict[str, str],
-        instruction: str,
-        workspace: str,
-        logs_dir: str | None = None,
-        pull: bool = True,
-        callback_url: str = "",
-        session_id: str = "",
-    ) -> AgentRunResult:
-        """Run an agent in a Docker container.
-
-        Pulls the image if needed, then runs the agent with proper
-        configuration. Returns a result object that can be iterated
-        to stream output lines.
-
-        Args:
-            image: Docker image URI
-            config: Agent-specific configuration dict (matches agent schema)
-            secrets: Secret values (API keys, tokens, etc.)
-            instruction: The task instruction/prompt for the agent
-            workspace: Host directory to mount as /workspace
-            logs_dir: Host directory for agent logs. If None, creates a temp dir.
-            pull: Whether to pull the image before running
-            callback_url: Full callback URL for Chronos (e.g., http://server/api/callback)
-            session_id: Chronos session ID for callbacks
-
-        Returns:
-            AgentRunResult that can be async-iterated for output lines.
-            Access result.logs_dir to get the host path where logs are stored.
-
-        Raises:
-            RuntimeError: If agent fails with non-zero exit code
-
-        Example:
-            result = AgentRunner.run(...)
-            async for line in result:
-                print(line)
-            print(f"Logs at: {result.logs_dir}")
-        """
-        return AgentRunResult(
-            image=image,
-            config=config,
-            secrets=secrets,
-            instruction=instruction,
-            workspace=workspace,
-            logs_dir=logs_dir,
-            pull=pull,
-            callback_url=callback_url,
-            session_id=session_id,
-        )
+from plato.agents.base import (
+    BaseAgent,
+    ConfigT,
+    get_agent,
+    get_registered_agents,
+    register_agent,
+)
+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.runner import run_agent
+from plato.agents.trajectory import (
+    SCHEMA_VERSION,
+    Agent,
+    FinalMetrics,
+    Metrics,
+    Observation,
+    ObservationResult,
+    Step,
+    ToolCall,
+    Trajectory,
+)

plato/agents/base.py

@@ -0,0 +1,145 @@
+"""Base agent class and registry for Plato agents."""
+
+from __future__ import annotations
+
+import json
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, ClassVar, Generic, TypeVar, get_args, get_origin
+
+from plato.agents.config import AgentConfig
+
+logger = logging.getLogger(__name__)
+
+# Global registry of agents
+_AGENT_REGISTRY: dict[str, type[BaseAgent]] = {}
+
+# Type variable for config
+ConfigT = TypeVar("ConfigT", bound=AgentConfig)
+
+
+def register_agent(name: str | None = None):
+    """Decorator to register an agent class.
+
+    Usage:
+        @register_agent("openhands")
+        class OpenHandsAgent(BaseAgent[OpenHandsConfig]):
+            ...
+    """
+
+    def decorator(cls: type[BaseAgent]) -> type[BaseAgent]:
+        agent_name = name or getattr(cls, "name", cls.__name__.lower().replace("agent", ""))
+        _AGENT_REGISTRY[agent_name] = cls
+        logger.debug(f"Registered agent: {agent_name} -> {cls.__name__}")
+        return cls
+
+    return decorator
+
+
+def get_registered_agents() -> dict[str, type[BaseAgent]]:
+    """Get all registered agents."""
+    return _AGENT_REGISTRY.copy()
+
+
+def get_agent(name: str) -> type[BaseAgent] | None:
+    """Get an agent by name."""
+    return _AGENT_REGISTRY.get(name)
+
+
+class BaseAgent(ABC, Generic[ConfigT]):
+    """Base class for Plato agents.
+
+    Subclass with a config type parameter for fully typed config access:
+
+        class OpenHandsConfig(AgentConfig):
+            model_name: str = "anthropic/claude-sonnet-4"
+            anthropic_api_key: Annotated[str | None, Secret(description="API key")] = None
+
+        @register_agent("openhands")
+        class OpenHandsAgent(BaseAgent[OpenHandsConfig]):
+            name = "openhands"
+            description = "OpenHands AI software engineer"
+
+            async def run(self, instruction: str) -> None:
+                # self.config is typed as OpenHandsConfig
+                model = self.config.model_name
+                ...
+    """
+
+    # Class attributes
+    name: ClassVar[str] = "base"
+    description: ClassVar[str] = ""
+
+    # Instance attributes
+    config: ConfigT
+
+    def __init__(self) -> None:
+        self.logger = logging.getLogger(f"plato.agents.{self.name}")
+
+    @classmethod
+    def get_config_class(cls) -> type[AgentConfig]:
+        """Get the config class from the generic parameter."""
+        for base in getattr(cls, "__orig_bases__", []):
+            origin = get_origin(base)
+            if origin is BaseAgent:
+                args = get_args(base)
+                if args and isinstance(args[0], type) and issubclass(args[0], AgentConfig):
+                    return args[0]
+        return AgentConfig
+
+    @classmethod
+    def get_version(cls) -> str:
+        """Get version from package metadata."""
+        import importlib.metadata
+
+        for pkg_name in [cls.__module__.split(".")[0], f"plato-agent-{cls.name}"]:
+            try:
+                return importlib.metadata.version(pkg_name)
+            except importlib.metadata.PackageNotFoundError:
+                continue
+        return "0.0.0"
+
+    @classmethod
+    def get_schema(cls) -> dict:
+        """Get full schema for the agent including config and build schemas."""
+        from plato.agents.build import BuildConfig
+
+        config_class = cls.get_config_class()
+        return {
+            "config": config_class.get_json_schema(),
+            "build": BuildConfig.get_json_schema(),
+        }
+
+    @abstractmethod
+    async def run(self, instruction: str) -> None:
+        """Run the agent with the given instruction.
+
+        This is the main entry point for agent execution. Implementations should:
+        1. Set up the environment using self.config
+        2. Execute the agent's core logic
+        3. Write trajectory to logs_dir if applicable
+
+        Args:
+            instruction: The task instruction/prompt for the agent.
+
+        Raises:
+            RuntimeError: If agent execution fails.
+        """
+        pass
+
+    async def write_trajectory(self, trajectory: dict[str, Any]) -> None:
+        """Write ATIF trajectory to the logs directory.
+
+        Args:
+            trajectory: ATIF-formatted trajectory dictionary.
+        """
+        logs_dir = Path(self.config.logs_dir)
+        agent_logs = logs_dir / "agent"
+        agent_logs.mkdir(parents=True, exist_ok=True)
+
+        trajectory_path = agent_logs / "trajectory.json"
+        with open(trajectory_path, "w") as f:
+            json.dump(trajectory, f, indent=2)
+
+        self.logger.info(f"Wrote trajectory to {trajectory_path}")