plato-sdk-v2 2.1.14__py3-none-any.whl → 2.1.16__py3-none-any.whl

plato/agents/runner.py

@@ -1,252 +1,90 @@
-"""Agent runner - discovers and runs Plato agents."""
+"""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 typing import Annotated
 
-import typer
-
-from plato.agents.callback import ChronosCallback
-
-app = typer.Typer(
-    name="plato-agent-runner",
-    help="Run Plato agents",
-    no_args_is_help=True,
-)
+from plato.agents.logging import log_event, span, upload_artifacts
 
 logger = logging.getLogger(__name__)
 
 
-def discover_agents() -> None:
-    """Discover and load installed agent packages via entry points.
-
-    Agent packages declare entry points in pyproject.toml:
-        [project.entry-points."plato.agents"]
-        openhands = "openhands_agent:OpenHandsAgent"
-
-    This function loads all such entry points, triggering registration.
-    """
-    import importlib.metadata
-
-    try:
-        eps = importlib.metadata.entry_points(group="plato.agents")
-    except TypeError:
-        # Python < 3.10 compatibility
-        eps = importlib.metadata.entry_points().get("plato.agents", [])
-
-    for ep in eps:
-        try:
-            ep.load()
-            logger.debug(f"Loaded agent: {ep.name}")
-        except Exception as e:
-            logger.warning(f"Failed to load agent '{ep.name}': {e}")
-
-
-@app.command()
-def run(
-    agent: Annotated[str, typer.Option("--agent", "-a", help="Agent name to run")],
-    instruction: Annotated[str, typer.Option("--instruction", "-i", help="Task instruction")],
-    config: Annotated[Path, typer.Option("--config", "-c", help="Path to config JSON file")],
-    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="Enable verbose logging")] = False,
+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 with the given instruction."""
-    # Setup logging
-    log_level = logging.DEBUG if verbose else logging.INFO
-    logging.basicConfig(
-        level=log_level,
-        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-    )
-
-    if not config.exists():
-        typer.echo(f"Error: Config file not found: {config}", err=True)
-        raise typer.Exit(1)
-
-    # Discover agents
-    discover_agents()
-
-    from plato.agents.base import get_agent, get_registered_agents
-
-    agent_cls = get_agent(agent)
-    if agent_cls is None:
-        available = list(get_registered_agents().keys())
-        typer.echo(f"Error: Agent '{agent}' not found. Available: {available}", err=True)
-        raise typer.Exit(1)
-
-    # Load config using the agent's typed config class
-    config_class = agent_cls.get_config_class()
-    agent_config = config_class.from_file(config)
-
-    try:
-        agent_instance = agent_cls()
-        agent_instance.config = agent_config
-        asyncio.run(agent_instance.run(instruction))
-    except Exception as e:
-        logger.exception(f"Agent execution failed: {e}")
-        raise typer.Exit(1)
-
-
-@app.command("list")
-def list_agents(
-    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="Enable verbose logging")] = False,
-) -> None:
-    """List available agents."""
-    if verbose:
-        logging.basicConfig(level=logging.DEBUG)
-
-    discover_agents()
-
-    from plato.agents.base import get_registered_agents
-
-    agents = get_registered_agents()
-    if not agents:
-        typer.echo("No agents found.")
-        return
-
-    typer.echo("Available agents:")
-    for name, cls in agents.items():
-        desc = getattr(cls, "description", "") or ""
-        version = cls.get_version()
-        typer.echo(f"  {name} (v{version}): {desc}")
-
-
-def main() -> None:
-    """CLI entry point."""
-    app()
-
-
-# =============================================================================
-# AgentRunner - Docker-based execution (existing functionality)
-# =============================================================================
-
-
-class AgentRunResult:
-    """Result of running an agent in Docker.
-
-    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.
+    """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]
 
-    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
-
-        self._callback = ChronosCallback(
-            callback_url=callback_url,
-            session_id=session_id,
-        )
-
-        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."""
-        return self._callback
+    async with span(agent_name, span_type="agent", source="agent") as agent_span:
+        agent_span.log(f"Starting agent: {agent_name} ({image})")
 
-    def __aiter__(self):
-        return self._stream()
-
-    async def _check_iptables_support(self) -> bool:
-        """Check if iptables is available."""
-        try:
-            proc = await asyncio.create_subprocess_exec(
-                "iptables",
-                "-L",
-                "-n",
-                stdout=asyncio.subprocess.DEVNULL,
-                stderr=asyncio.subprocess.DEVNULL,
-            )
-            await proc.wait()
-            return proc.returncode == 0
-        except (FileNotFoundError, PermissionError):
-            return False
-
-    async def _stream(self):
-        """Stream output from the agent."""
-        import json
-        import os
-        import tempfile
-
-        log_buffer: list[dict] = []
-
-        async def flush_logs():
-            if not log_buffer or not self._callback.enabled:
-                return
-            await self._callback.push_logs(log_buffer.copy())
-            log_buffer.clear()
-
-        agent_name = self._image.split("/")[-1].split(":")[0]
-        await self._callback.push_log(f"Starting agent: {agent_name} ({self._image})")
-
-        if self._pull:
+        # Pull image if requested
+        if pull:
+            agent_span.log(f"Pulling image: {image}")
             pull_proc = await asyncio.create_subprocess_exec(
                 "docker",
                 "pull",
-                self._image,
+                image,
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.STDOUT,
             )
-            assert pull_proc.stdout is not None
-            while True:
-                line = await pull_proc.stdout.readline()
-                if not line:
-                    break
-                yield f"[pull] {line.decode().rstrip()}"
             await pull_proc.wait()
 
-        agent_logs_subdir = os.path.join(self._logs_dir, "agent")
-        os.makedirs(agent_logs_subdir, exist_ok=True)
-
+        # Setup
+        os.makedirs(os.path.join(logs_dir, "agent"), exist_ok=True)
         config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
-        json.dump(self._config, config_file)
+        json.dump(config, config_file)
         config_file.close()
 
-        agent_failed = False
-        error_message = ""
-
         try:
-            use_host_network = not await self._check_iptables_support()
-            if use_host_network:
-                yield "[info] iptables not available, using host network mode"
-
+            # Build docker command
             docker_cmd = ["docker", "run", "--rm"]
 
-            if use_host_network:
+            # 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"{self._workspace}:/workspace",
+                    f"{workspace}:/workspace",
                     "-v",
-                    f"{self._logs_dir}:/logs",
+                    f"{logs_dir}:/logs",
                     "-v",
                     f"{config_file.name}:/config.json:ro",
                     "-w",
@@ -254,121 +92,59 @@ class AgentRunResult:
                 ]
             )
 
-            for key, value in self._secrets.items():
+            for key, value in secrets.items():
                 docker_cmd.extend(["-e", f"{key.upper()}={value}"])
 
-            docker_cmd.append(self._image)
-            docker_cmd.extend(["--instruction", self._instruction])
+            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,
             )
-            assert process.stdout is not None
 
+            # Stream output line by line
+            assert process.stdout is not None
             while True:
                 line = await process.stdout.readline()
                 if not line:
                     break
-                decoded = line.decode().rstrip()
-                yield decoded
-
-                if self._callback.enabled:
-                    log_buffer.append({"level": "info", "message": decoded})
-                    if len(log_buffer) >= 10:
-                        await flush_logs()
+                logger.info(f"[agent] {line.decode().rstrip()}")
 
             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:
-            os.unlink(config_file.name)
-            await flush_logs()
-
-            if self._callback.enabled:
-                logger.info(f"Callback enabled, uploading artifacts from {self._logs_dir}")
-                if agent_failed:
-                    await self._callback.push_log(error_message, level="error")
-                else:
-                    await self._callback.push_log("Agent completed successfully")
-                result = await self._callback.upload_artifacts(
-                    logs_dir=self._logs_dir,
-                    agent_image=self._image,
-                )
-                logger.info(f"Artifact upload result: {result}")
-            else:
-                logger.info(
-                    f"Callback not enabled (url={self._callback.callback_url}, session={self._callback.session_id})"
-                )
-
-        if agent_failed:
-            raise RuntimeError(error_message)
-
+                raise RuntimeError(f"Agent failed with exit code {process.returncode}")
 
-class AgentRunner:
-    """Run agents in Docker containers.
-
-    This provides isolated execution of agents. For direct execution,
-    use the plato-agent-runner CLI or instantiate the agent directly.
-
-    Example:
-        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",
-            workspace="/path/to/repo",
-        ):
-            print(line)
-    """
-
-    @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.
-
-        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
-            callback_url: Chronos callback URL
-            session_id: Chronos session ID
-
-        Returns:
-            AgentRunResult that can be async-iterated for output.
-        """
-        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,
-        )
+            agent_span.log("Agent completed successfully")
 
+        finally:
+            os.unlink(config_file.name)
 
-if __name__ == "__main__":
-    main()
+            # 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)