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

plato/agents/runner.py

@@ -10,7 +10,10 @@ import platform
 import tempfile
 from pathlib import Path
 
-from plato.agents.logging import log_event, span, upload_artifacts
+from opentelemetry import trace
+
+from plato.agents.artifacts import upload_artifacts
+from plato.agents.otel import get_tracer
 
 logger = logging.getLogger(__name__)
 
@@ -38,20 +41,32 @@ async def run_agent(
     logs_dir = logs_dir or tempfile.mkdtemp(prefix="agent_logs_")
     agent_name = image.split("/")[-1].split(":")[0]
 
-    async with span(agent_name, span_type="agent", source="agent") as agent_span:
-        agent_span.log(f"Starting agent: {agent_name} ({image})")
+    # Get session info from environment variables
+    session_id = os.environ.get("SESSION_ID")
+    otel_url = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
+    upload_url = os.environ.get("UPLOAD_URL")
+
+    tracer = get_tracer("plato.agent")
+
+    with tracer.start_as_current_span(agent_name) as agent_span:
+        agent_span.set_attribute("span.type", "agent")
+        agent_span.set_attribute("source", "agent")
+        agent_span.set_attribute("image", image)
+        agent_span.set_attribute("content", f"Starting agent: {agent_name}")
 
         # Pull image if requested
         if pull:
-            agent_span.log(f"Pulling image: {image}")
-            pull_proc = await asyncio.create_subprocess_exec(
-                "docker",
-                "pull",
-                image,
-                stdout=asyncio.subprocess.PIPE,
-                stderr=asyncio.subprocess.STDOUT,
-            )
-            await pull_proc.wait()
+            with tracer.start_as_current_span("docker_pull") as pull_span:
+                pull_span.set_attribute("span.type", "docker_pull")
+                pull_span.set_attribute("image", image)
+                pull_proc = await asyncio.create_subprocess_exec(
+                    "docker",
+                    "pull",
+                    image,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.STDOUT,
+                )
+                await pull_proc.wait()
 
         # Setup
         os.makedirs(os.path.join(logs_dir, "agent"), exist_ok=True)
@@ -63,9 +78,7 @@ async def run_agent(
             # Build docker command
             docker_cmd = ["docker", "run", "--rm"]
 
-            # Determine if we need host networking:
-            # - Required on Linux without iptables for connectivity
-            # - Skip on macOS where --network=host doesn't work properly
+            # Determine if we need host networking
             use_host_network = False
             is_macos = platform.system() == "Darwin"
 
@@ -103,68 +116,200 @@ async def run_agent(
                 ]
             )
 
+            # Pass session info to agent
+            if otel_url:
+                docker_cmd.extend(["-e", f"OTEL_EXPORTER_OTLP_ENDPOINT={otel_url}"])
+                # Use JSON protocol (not protobuf) for OTLP exports
+                docker_cmd.extend(["-e", "OTEL_EXPORTER_OTLP_PROTOCOL=http/json"])
+            if session_id:
+                docker_cmd.extend(["-e", f"SESSION_ID={session_id}"])
+            if upload_url:
+                docker_cmd.extend(["-e", f"UPLOAD_URL={upload_url}"])
+
+            # Pass trace context to agent for parent linking
+            current_span = trace.get_current_span()
+            span_context = current_span.get_span_context()
+            if span_context.is_valid:
+                trace_id = format(span_context.trace_id, "032x")
+                span_id = format(span_context.span_id, "016x")
+                docker_cmd.extend(
+                    [
+                        "-e",
+                        f"OTEL_TRACE_ID={trace_id}",
+                        "-e",
+                        f"OTEL_PARENT_SPAN_ID={span_id}",
+                    ]
+                )
+
             for key, value in secrets.items():
                 docker_cmd.extend(["-e", f"{key.upper()}={value}"])
 
             docker_cmd.append(image)
 
-            # Pass instruction via CLI arg (agents expect --instruction flag)
+            # Pass instruction via CLI arg
             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,
-            )
+            with tracer.start_as_current_span("agent_execution") as exec_span:
+                exec_span.set_attribute("span.type", "agent_execution")
+                exec_span.set_attribute("content", f"Running {agent_name}")
 
-            # Stream output line by line, collecting for error reporting
-            output_lines: list[str] = []
-            assert process.stdout is not None
-            while True:
-                line = await process.stdout.readline()
-                if not line:
-                    break
-                decoded_line = line.decode().rstrip()
-                output_lines.append(decoded_line)
-                logger.info(f"[agent] {decoded_line}")
-
-            await process.wait()
-
-            if process.returncode != 0:
-                # Get last N lines of output for error context
-                error_context = "\n".join(output_lines[-50:]) if output_lines else "No output captured"
-
-                # Log error event with container output
-                await log_event(
-                    span_type="error",
-                    content=f"Agent failed with exit code {process.returncode}",
-                    source="agent",
-                    extra={
-                        "exit_code": process.returncode,
-                        "image": image,
-                        "agent_name": agent_name,
-                        "output": error_context,
-                        "output_line_count": len(output_lines),
-                    },
+                process = await asyncio.create_subprocess_exec(
+                    *docker_cmd,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.STDOUT,
                 )
 
-                agent_span.set_extra(
-                    {
-                        "error": True,
-                        "exit_code": process.returncode,
-                        "output": error_context,
-                    }
-                )
+                # Stream output line by line
+                output_lines: list[str] = []
+                turn_count = 0
+                assert process.stdout is not None
+                while True:
+                    line = await process.stdout.readline()
+                    if not line:
+                        break
+                    decoded_line = line.decode().rstrip()
+                    output_lines.append(decoded_line)
+
+                    # Try to parse JSON output from agent for structured trajectory spans
+                    try:
+                        data = json.loads(decoded_line)
+                        event_type = data.get("type", "")
+
+                        if event_type == "assistant":
+                            # Agent response - create a turn span
+                            turn_count += 1
+                            msg = data.get("message", {})
+                            content_items = msg.get("content", [])
+
+                            # Extract text and tool calls with full details
+                            text_parts = []
+                            tool_calls = []
+                            for item in content_items:
+                                if item.get("type") == "text":
+                                    text_parts.append(item.get("text", "")[:2000])
+                                elif item.get("type") == "tool_use":
+                                    tool_input = item.get("input", {})
+                                    # Truncate large inputs
+                                    input_str = json.dumps(tool_input) if tool_input else ""
+                                    if len(input_str) > 2000:
+                                        input_str = input_str[:2000] + "..."
+                                    tool_calls.append(
+                                        {
+                                            "tool": item.get("name"),
+                                            "id": item.get("id"),
+                                            "input": input_str,
+                                        }
+                                    )
+
+                            with tracer.start_as_current_span(f"turn_{turn_count}") as turn_span:
+                                turn_span.set_attribute("span.type", "agent_turn")
+                                turn_span.set_attribute("source", "agent")
+                                turn_span.set_attribute("turn_number", turn_count)
+                                turn_span.set_attribute("model", msg.get("model", "unknown"))
 
-                raise RuntimeError(f"Agent failed with exit code {process.returncode}")
+                                if text_parts:
+                                    turn_span.set_attribute("content", "\n".join(text_parts)[:4000])
+                                if tool_calls:
+                                    turn_span.set_attribute("tool_calls", json.dumps(tool_calls))
+                                    # If no text content, show tool calls summary
+                                    if not text_parts:
+                                        turn_span.set_attribute(
+                                            "content", f"Tool calls: {', '.join(t['tool'] for t in tool_calls)}"
+                                        )
 
-            agent_span.log("Agent completed successfully")
+                                # Usage info
+                                usage = msg.get("usage", {})
+                                if usage:
+                                    turn_span.set_attribute("input_tokens", usage.get("input_tokens", 0))
+                                    turn_span.set_attribute("output_tokens", usage.get("output_tokens", 0))
+
+                        elif event_type == "user":
+                            # Tool result
+                            tool_results = data.get("message", {}).get("content", [])
+                            for result in tool_results:
+                                if result.get("type") == "tool_result":
+                                    tool_id = result.get("tool_use_id", "")
+                                    content = result.get("content", "")
+                                    # Handle content that might be a list of content blocks
+                                    if isinstance(content, list):
+                                        text_parts = []
+                                        for item in content:
+                                            if isinstance(item, dict) and item.get("type") == "text":
+                                                text_parts.append(item.get("text", ""))
+                                            elif isinstance(item, str):
+                                                text_parts.append(item)
+                                        content = "\n".join(text_parts)
+                                    if isinstance(content, str):
+                                        content = content[:2000]  # Truncate large results
+                                    with tracer.start_as_current_span("tool_result") as tr_span:
+                                        tr_span.set_attribute("span.type", "tool_result")
+                                        tr_span.set_attribute("source", "agent")
+                                        tr_span.set_attribute("tool_use_id", tool_id)
+                                        tr_span.set_attribute("content", f"Tool result for {tool_id}")
+                                        tr_span.set_attribute("result", content if content else "")
+
+                        elif event_type == "result":
+                            # Final result
+                            result_text = data.get("result", "")[:1000]
+                            is_error = data.get("is_error", False)
+                            duration_ms = data.get("duration_ms", 0)
+                            total_cost = data.get("total_cost_usd", 0)
+
+                            with tracer.start_as_current_span("agent_result") as res_span:
+                                res_span.set_attribute("span.type", "agent_result")
+                                res_span.set_attribute("source", "agent")
+                                res_span.set_attribute("content", result_text if result_text else "Agent completed")
+                                res_span.set_attribute("is_error", is_error)
+                                res_span.set_attribute("duration_ms", duration_ms)
+                                res_span.set_attribute("total_cost_usd", total_cost)
+                                res_span.set_attribute("num_turns", data.get("num_turns", turn_count))
+
+                        elif event_type == "system" and data.get("subtype") == "init":
+                            # Agent initialization
+                            with tracer.start_as_current_span("agent_init") as init_span:
+                                init_span.set_attribute("span.type", "agent_init")
+                                init_span.set_attribute("source", "agent")
+                                init_span.set_attribute("model", data.get("model", "unknown"))
+                                init_span.set_attribute("tools", json.dumps(data.get("tools", [])))
+                                init_span.set_attribute("content", f"Agent initialized: {data.get('model', 'unknown')}")
+
+                        else:
+                            # Other output - just log it without creating a span
+                            logger.debug(f"[agent] {decoded_line}")
+                            continue
+
+                    except json.JSONDecodeError:
+                        # Not JSON - just log it
+                        logger.info(f"[agent] {decoded_line}")
+
+                await process.wait()
+
+                exit_code = process.returncode or 0
+                if exit_code != 0:
+                    error_context = "\n".join(output_lines[-50:]) if output_lines else "No output captured"
+
+                    exec_span.set_attribute("error", True)
+                    exec_span.set_attribute("exit_code", exit_code)
+                    exec_span.add_event(
+                        "agent_error",
+                        {
+                            "exit_code": exit_code,
+                            "output": error_context[:4000],
+                        },
+                    )
+
+                    agent_span.set_attribute("error", True)
+                    agent_span.set_attribute("exit_code", exit_code)
+
+                    raise RuntimeError(f"Agent failed with exit code {exit_code}")
+
+                exec_span.set_attribute("success", True)
 
         finally:
             os.unlink(config_file.name)
 
-            # Load trajectory and add to span
+            # Load trajectory and log as event
             trajectory_path = Path(logs_dir) / "agent" / "trajectory.json"
             if trajectory_path.exists():
                 try:
@@ -178,14 +323,24 @@ async def run_agent(
                         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",
-                        )
+                        # Log trajectory as span event
+                        with tracer.start_as_current_span("trajectory") as traj_span:
+                            traj_span.set_attribute("span.type", "trajectory")
+                            traj_span.set_attribute("log_type", "atif")
+                            traj_span.set_attribute("source", "agent")
+                            # Store trajectory in span (truncated for OTel limits)
+                            traj_json = json.dumps(trajectory)
+                            if len(traj_json) > 10000:
+                                traj_span.set_attribute("trajectory_truncated", True)
+                                traj_span.set_attribute("trajectory_size", len(traj_json))
+                            else:
+                                traj_span.set_attribute("trajectory", traj_json)
                 except Exception as e:
                     logger.warning(f"Failed to load trajectory: {e}")
 
-            await upload_artifacts(logs_dir)
+            # Upload artifacts if we have upload URL configured
+            if upload_url:
+                await upload_artifacts(upload_url, logs_dir)
+
+        agent_span.set_attribute("success", True)
+        agent_span.set_attribute("content", f"Agent {agent_name} completed successfully")

plato/worlds/README.md

@@ -28,7 +28,8 @@ Create a JSON config file:
   },
   "git_token": "ghp_...",
   "session_id": "local-test-001",
-  "callback_url": ""
+  "otel_url": "",
+  "upload_url": ""
 }
 ```
 

plato/worlds/base.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import os
 import subprocess
 from abc import ABC, abstractmethod
 from pathlib import Path
@@ -16,12 +17,14 @@ if TYPE_CHECKING:
     from plato.v2.async_.environment import Environment
     from plato.v2.async_.session import Session
 
-from plato.agents.logging import init_logging as _init_chronos_logging
-from plato.agents.logging import log_event as _log_event
-from plato.agents.logging import reset_logging as _reset_chronos_logging
-from plato.agents.logging import span as _span
-from plato.agents.logging import upload_artifact as _upload_artifact
-from plato.agents.logging import upload_checkpoint as _upload_checkpoint
+from plato.agents.artifacts import (
+    upload_artifact as _upload_artifact_raw,
+)
+from plato.agents.otel import (
+    get_tracer,
+    init_tracing,
+    shutdown_tracing,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -111,6 +114,7 @@ class BaseWorld(ABC, Generic[ConfigT]):
         self._step_count: int = 0
         self.plato_session = None
         self._current_step_id: str | None = None
+        self._session_id: str | None = None
 
     @classmethod
     def get_config_class(cls) -> type[RunConfig]:
@@ -390,17 +394,39 @@ class BaseWorld(ABC, Generic[ConfigT]):
             self.logger.warning(f"Failed to create state bundle: {e.stderr}")
             return None
 
-    async def _create_and_upload_checkpoint(self) -> dict[str, Any] | None:
+    async def _upload_artifact(
+        self,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> bool:
+        """Upload an artifact directly to S3.
+
+        Args:
+            data: Raw bytes of the artifact
+            content_type: MIME type of the content
+
+        Returns:
+            True if successful, False otherwise
+        """
+        if not self.config.upload_url:
+            self.logger.warning("Cannot upload artifact: upload_url not set")
+            return False
+        return await _upload_artifact_raw(
+            upload_url=self.config.upload_url,
+            data=data,
+            content_type=content_type,
+        )
+
+    async def _create_and_upload_checkpoint(self) -> dict[str, str] | None:
         """Create a full checkpoint including env snapshots and state bundle.
 
         This method:
         1. Commits any pending state changes
         2. Creates env snapshots using snapshot_store
-        3. Creates and uploads state bundle as an artifact
-        4. Calls the checkpoint endpoint with all data
+        3. Creates and uploads state bundle to S3
 
         Returns:
-            Checkpoint result dict if successful, None otherwise.
+            Dict mapping env alias to artifact_id if successful, None otherwise
         """
         # Commit state changes first
         self._commit_state(f"Checkpoint at step {self._step_count}")
@@ -411,35 +437,20 @@ class BaseWorld(ABC, Generic[ConfigT]):
             env_snapshots = {}
 
         # Create and upload state bundle
-        state_artifact_id: str | None = None
         if self.config.state.enabled:
             bundle_data = self._create_state_bundle()
             if bundle_data:
-                result = await _upload_artifact(
+                success = await self._upload_artifact(
                     data=bundle_data,
-                    artifact_type="state",
-                    filename=f"state_step_{self._step_count}.bundle",
-                    extra={
-                        "step_number": self._step_count,
-                        "state_path": self.config.state.path,
-                    },
+                    content_type="application/octet-stream",
                 )
-                if result:
-                    state_artifact_id = result.get("artifact_id")
-                    self.logger.info(f"Uploaded state artifact: {state_artifact_id}")
-
-        # Upload checkpoint with all data
-        checkpoint_result = await _upload_checkpoint(
-            step_number=self._step_count,
-            env_snapshots=env_snapshots,
-            state_artifact_id=state_artifact_id,
-            extra={
-                "world_name": self.name,
-                "world_version": self.get_version(),
-            },
-        )
+                if success:
+                    self.logger.info(f"Uploaded state bundle at step {self._step_count}")
+                else:
+                    self.logger.warning(f"Failed to upload state bundle at step {self._step_count}")
+                    return None
 
-        return checkpoint_result
+        return env_snapshots
 
     def get_env(self, alias: str) -> Environment | None:
         """Get an environment by alias.
@@ -630,64 +641,98 @@ The following services are available for your use:
         # Initialize state directory (creates git repo if needed)
         self._init_state_directory()
 
-        # Initialize the logging singleton for agents to use
-        if config.callback_url and config.session_id:
-            _init_chronos_logging(
-                callback_url=config.callback_url,
-                session_id=config.session_id,
-            )
+        # Initialize OTel tracing and session info for artifact uploads
+        if config.session_id:
+            self._session_id = config.session_id
+
+            # Set environment variables for agent runners
+            os.environ["SESSION_ID"] = config.session_id
+            if config.otel_url:
+                os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = config.otel_url
+                # Use JSON protocol (not protobuf) for the OTLP exporter
+                os.environ["OTEL_EXPORTER_OTLP_PROTOCOL"] = "http/json"
+            if config.upload_url:
+                os.environ["UPLOAD_URL"] = config.upload_url
+
+            # Initialize OTel tracing if otel_url is provided
+            print(f"[World] OTel URL from config: {config.otel_url!r}")
+            if config.otel_url:
+                init_tracing(
+                    service_name=f"world-{self.name}",
+                    session_id=config.session_id,
+                    otlp_endpoint=config.otel_url,
+                )
+            else:
+                print("[World] No otel_url in config - OTel tracing disabled")
 
         # Connect to Plato session if configured (for heartbeats)
         await self._connect_plato_session()
 
+        # Get tracer for spans
+        tracer = get_tracer("plato.world")
+
         # Log session start
-        await _log_event(
-            span_type="session_start",
-            content=f"World '{self.name}' started",
-            source="world",
-            extra={"world_name": self.name, "world_version": self.get_version()},
-        )
+        with tracer.start_as_current_span("session_start") as span:
+            span.set_attribute("span.type", "session_start")
+            span.set_attribute("source", "world")
+            span.set_attribute("world_name", self.name)
+            span.set_attribute("world_version", self.get_version())
+            span.set_attribute("content", f"World '{self.name}' started")
 
         try:
-            # Execute reset with automatic span tracking
-            async with _span("reset", span_type="reset", source="world") as reset_span:
-                reset_span.log(f"Resetting world '{self.name}'")
+            # Execute reset with OTel span
+            with tracer.start_as_current_span("reset") as reset_span:
+                reset_span.set_attribute("span.type", "reset")
+                reset_span.set_attribute("source", "world")
+                reset_span.set_attribute("content", f"Resetting world '{self.name}'")
                 obs = await self.reset()
-                reset_span.set_extra({"observation": obs.model_dump() if hasattr(obs, "model_dump") else str(obs)})
+                obs_data = obs.model_dump() if hasattr(obs, "model_dump") else str(obs)
+                reset_span.set_attribute("observation", str(obs_data)[:1000])  # Truncate for OTel
             self.logger.info(f"World reset complete: {obs}")
 
             while True:
                 self._step_count += 1
 
-                # Execute step with automatic span tracking
-                # The span automatically sets itself as the current parent,
-                # so agent trajectories will nest under this step
-                async with _span(
-                    f"step_{self._step_count}",
-                    span_type="step",
-                    source="world",
-                ) as step_span:
-                    self._current_step_id = step_span.event_id
-                    step_span.log(f"Step {self._step_count} started")
+                # Execute step with OTel span
+                with tracer.start_as_current_span(f"step_{self._step_count}") as step_span:
+                    step_span.set_attribute("span.type", "step")
+                    step_span.set_attribute("source", "world")
+                    step_span.set_attribute("step_number", self._step_count)
+                    step_span.set_attribute("content", f"Step {self._step_count} started")
+
+                    # Store span context for nested agent spans
+
+                    self._current_step_id = format(step_span.get_span_context().span_id, "016x")
+
                     result = await self.step()
-                    step_span.set_extra(
-                        {
-                            "done": result.done,
-                            "observation": result.observation.model_dump()
-                            if hasattr(result.observation, "model_dump")
-                            else str(result.observation),
-                            "info": result.info,
-                        }
+
+                    step_span.set_attribute("done", result.done)
+                    obs_data = (
+                        result.observation.model_dump()
+                        if hasattr(result.observation, "model_dump")
+                        else str(result.observation)
                     )
+                    step_span.set_attribute("observation", str(obs_data)[:1000])
 
                 self.logger.info(f"Step {self._step_count}: done={result.done}")
 
                 # Create checkpoint if enabled and interval matches
-                # Note: The checkpoint event is created by the callback endpoint,
-                # so we don't need a span wrapper here (would create duplicates)
                 if self.config.checkpoint.enabled and self._step_count % self.config.checkpoint.interval == 0:
                     self.logger.info(f"Creating checkpoint after step {self._step_count}")
-                    await self._create_and_upload_checkpoint()
+                    env_snapshots = await self._create_and_upload_checkpoint()
+
+                    # Emit checkpoint span for UI visibility
+                    with tracer.start_as_current_span(f"checkpoint_{self._step_count}") as checkpoint_span:
+                        checkpoint_span.set_attribute("span.type", "checkpoint")
+                        checkpoint_span.set_attribute("source", "world")
+                        checkpoint_span.set_attribute("step_number", self._step_count)
+                        checkpoint_span.set_attribute("content", f"Checkpoint created at step {self._step_count}")
+                        if env_snapshots:
+                            # Serialize env_snapshots for OTel attribute
+                            import json
+
+                            checkpoint_span.set_attribute("env_snapshots", json.dumps(env_snapshots))
+                        checkpoint_span.set_attribute("success", env_snapshots is not None)
 
                 if result.done:
                     break
@@ -697,14 +742,14 @@ The following services are available for your use:
             await self._disconnect_plato_session()
 
             # Log session end
-            await _log_event(
-                span_type="session_end",
-                content=f"World '{self.name}' completed after {self._step_count} steps",
-                source="world",
-                extra={"total_steps": self._step_count},
-            )
-
-            # Reset the logging singleton
-            _reset_chronos_logging()
+            with tracer.start_as_current_span("session_end") as span:
+                span.set_attribute("span.type", "session_end")
+                span.set_attribute("source", "world")
+                span.set_attribute("total_steps", self._step_count)
+                span.set_attribute("content", f"World '{self.name}' completed after {self._step_count} steps")
+
+            # Shutdown OTel tracing and clear session info
+            shutdown_tracing()
+            self._session_id = None
 
             self.logger.info(f"World '{self.name}' completed after {self._step_count} steps")