droidrun 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

droidrun/agent/utils/trajectory.py

@@ -0,0 +1,184 @@
+"""
+Trajectory utilities for DroidRun agents.
+
+This module provides helper functions for working with agent trajectories,
+including saving, loading, and analyzing them.
+"""
+
+import json
+import logging
+import os
+import time
+from typing import Dict, List, Any
+from PIL import Image
+import io
+from llama_index.core.workflow import Event
+
+logger = logging.getLogger("droidrun")
+
+class Trajectory:
+
+    def __init__(self):
+        """Initializes an empty trajectory class."""
+        self.events: List[Event] = [] 
+        self.screenshots: List[bytes] = [] 
+
+
+    def create_screenshot_gif(self, output_path: str, duration: int = 1000) -> str:
+        """
+        Create a GIF from a list of screenshots.
+        
+        Args:
+            output_path: Base path for the GIF (without extension)
+            duration: Duration for each frame in milliseconds
+        
+        Returns:
+            Path to the created GIF file
+        """
+        if len(self.screenshots) == 0:
+            return None
+            
+        images = []
+        for screenshot in self.screenshots:
+            img_data = screenshot
+            img = Image.open(io.BytesIO(img_data))
+            images.append(img)
+        
+        # Save as GIF
+        gif_path = f"{output_path}.gif"
+        images[0].save(
+            gif_path,
+            save_all=True,
+            append_images=images[1:],
+            duration=duration,
+            loop=0
+        )
+        
+        return gif_path
+
+    def save_trajectory(
+        self,
+        directory: str = "trajectories",
+    ) -> str:
+        """
+        Save trajectory steps to a JSON file and create a GIF of screenshots if available.
+        
+        Args:
+            directory: Directory to save the trajectory files
+        
+        Returns:
+            Path to the saved trajectory file
+        """
+        os.makedirs(directory, exist_ok=True)
+        
+        timestamp = time.strftime("%Y%m%d_%H%M%S")
+        base_path = os.path.join(directory, f"trajectory_{timestamp}")
+        
+        def make_serializable(obj):
+            """Recursively make objects JSON serializable."""
+            if hasattr(obj, "__class__") and obj.__class__.__name__ == "ChatMessage":
+                # Extract the text content from the ChatMessage
+                if hasattr(obj, "content") and obj.content is not None:
+                    return {"role": obj.role.value, "content": obj.content}
+                # If content is not available, try extracting from blocks
+                elif hasattr(obj, "blocks") and obj.blocks:
+                    text_content = ""
+                    for block in obj.blocks:
+                        if hasattr(block, "text"):
+                            text_content += block.text
+                    return {"role": obj.role.value, "content": text_content}
+                else:
+                    return str(obj)
+            elif isinstance(obj, dict):
+                return {k: make_serializable(v) for k, v in obj.items()}
+            elif isinstance(obj, list):
+                return [make_serializable(item) for item in obj]
+            elif hasattr(obj, "__dict__"):
+                # Handle other custom objects by converting to dict
+                return {k: make_serializable(v) for k, v in obj.__dict__.items() 
+                       if not k.startswith('_')}
+            else:
+                return obj
+        
+        serializable_events = []
+        for event in self.events:
+            event_dict = {
+                "type": event.__class__.__name__,
+                **{k: make_serializable(v) for k, v in event.__dict__.items() 
+                   if not k.startswith('_')}
+            }
+            serializable_events.append(event_dict)
+        
+        json_path = f"{base_path}.json"
+        with open(json_path, "w") as f:
+            json.dump(serializable_events, f, indent=2)
+
+        self.create_screenshot_gif(base_path)
+
+        return json_path
+
+    def get_trajectory_statistics(trajectory_data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Get statistics about a trajectory.
+        
+        Args:
+            trajectory_data: The trajectory data dictionary
+            
+        Returns:
+            Dictionary with statistics about the trajectory
+        """
+        trajectory_steps = trajectory_data.get("trajectory_steps", [])
+        
+        # Count different types of steps
+        step_types = {}
+        for step in trajectory_steps:
+            step_type = step.get("type", "unknown")
+            step_types[step_type] = step_types.get(step_type, 0) + 1
+        
+        # Count planning vs execution steps
+        planning_steps = sum(count for step_type, count in step_types.items() 
+                            if step_type.startswith("planner_"))
+        execution_steps = sum(count for step_type, count in step_types.items() 
+                            if step_type.startswith("codeact_"))
+        
+        # Count successful vs failed executions
+        successful_executions = sum(1 for step in trajectory_steps 
+                                if step.get("type") == "codeact_execution" 
+                                and step.get("success", False))
+        failed_executions = sum(1 for step in trajectory_steps 
+                            if step.get("type") == "codeact_execution" 
+                            and not step.get("success", True))
+        
+        # Return statistics
+        return {
+            "total_steps": len(trajectory_steps),
+            "step_types": step_types,
+            "planning_steps": planning_steps,
+            "execution_steps": execution_steps,
+            "successful_executions": successful_executions,
+            "failed_executions": failed_executions,
+            "goal_achieved": trajectory_data.get("success", False)
+        }
+
+    def print_trajectory_summary(self, trajectory_data: Dict[str, Any]) -> None:
+        """
+        Print a summary of a trajectory.
+        
+        Args:
+            trajectory_data: The trajectory data dictionary
+        """
+        stats = self.get_trajectory_statistics(trajectory_data)
+        
+        print("=== Trajectory Summary ===")
+        print(f"Goal: {trajectory_data.get('goal', 'Unknown')}")
+        print(f"Success: {trajectory_data.get('success', False)}")
+        print(f"Reason: {trajectory_data.get('reason', 'Unknown')}")
+        print(f"Total steps: {stats['total_steps']}")
+        print("Step breakdown:")
+        for step_type, count in stats['step_types'].items():
+            print(f"  - {step_type}: {count}")
+        print(f"Planning steps: {stats['planning_steps']}")
+        print(f"Execution steps: {stats['execution_steps']}")
+        print(f"Successful executions: {stats['successful_executions']}")
+        print(f"Failed executions: {stats['failed_executions']}")
+        print("==========================")

droidrun/cli/__init__.py

@@ -4,6 +4,6 @@ DroidRun CLI Module.
 This module provides command-line interfaces for interacting with Android devices.
 """
 
-from .main import cli
+from droidrun.cli.main import cli
 
 __all__ = ["cli"] 

droidrun/cli/logs.py

@@ -0,0 +1,283 @@
+import logging
+from rich.layout import Layout
+from rich.panel import Panel
+from rich.spinner import Spinner
+from rich.console import Console
+from rich.live import Live
+from typing import List
+
+from droidrun.agent.common.events import ScreenshotEvent
+from droidrun.agent.planner.events import (
+    PlanInputEvent,
+    PlanThinkingEvent,
+    PlanCreatedEvent,
+)
+from droidrun.agent.codeact.events import (
+    TaskInputEvent,
+    TaskThinkingEvent,
+    TaskExecutionEvent,
+    TaskExecutionResultEvent,
+    TaskEndEvent,
+)
+from droidrun.agent.droid.events import (
+    CodeActExecuteEvent,
+    CodeActResultEvent,
+    ReasoningLogicEvent,
+    TaskRunnerEvent,
+    FinalizeEvent,
+)
+
+
+class LogHandler(logging.Handler):
+    def __init__(self, goal: str, current_step: str = "Initializing..."):
+        super().__init__()
+
+        self.goal = goal
+        self.current_step = current_step
+        self.is_completed = False
+        self.is_success = False
+        self.spinner = Spinner("dots")
+        self.console = Console()
+        self.layout = self._create_layout()
+        self.logs: List[str] = []
+
+    def emit(self, record):
+        msg = self.format(record)
+        lines = msg.splitlines()
+
+        for line in lines:
+            self.logs.append(line)
+            # Optionally, limit the log list size
+            if len(self.logs) > 100:
+                self.logs.pop(0)
+
+        self.rerender()
+
+    def render(self):
+        return Live(self.layout, refresh_per_second=4, console=self.console)
+
+    def rerender(self):
+        self._update_layout(
+            self.layout,
+            self.logs,
+            self.current_step,
+            self.goal,
+            self.is_completed,
+            self.is_success,
+        )
+
+    def update_step(self, step: str):
+        self.current_step = step
+        self.rerender()
+
+    def _create_layout(self):
+        """Create a layout with logs at top and status at bottom"""
+        layout = Layout()
+        layout.split(
+            Layout(name="logs"),
+            Layout(name="goal", size=3),
+            Layout(name="status", size=3),
+        )
+        return layout
+
+    def _update_layout(
+        self,
+        layout: Layout,
+        log_list: List[str],
+        step_message: str,
+        goal: str = None,
+        completed: bool = False,
+        success: bool = False,
+    ):
+        """Update the layout with current logs and step information"""
+        from rich.text import Text
+        import shutil
+
+        # Cache terminal size to avoid frequent recalculation
+        try:
+            terminal_height = shutil.get_terminal_size().lines
+        except:
+            terminal_height = 24  # fallback
+
+        # Reserve space for panels and borders (more conservative estimate)
+        other_components_height = 10  # goal panel + status panel + borders + padding
+        available_log_lines = max(8, terminal_height - other_components_height)
+
+        # Only show recent logs, but ensure we don't flicker
+        visible_logs = (
+            log_list[-available_log_lines:]
+            if len(log_list) > available_log_lines
+            else log_list
+        )
+
+        # Ensure we always have some content to prevent panel collapse
+        if not visible_logs:
+            visible_logs = ["Initializing..."]
+
+        log_content = "\n".join(visible_logs)
+
+        layout["logs"].update(
+            Panel(
+                log_content,
+                title=f"Activity Log ({len(log_list)} entries)",
+                border_style="blue",
+                title_align="left",
+                padding=(0, 1),
+                height=available_log_lines + 2,
+            )
+        )
+
+        if goal:
+            goal_text = Text(goal, style="bold")
+            layout["goal"].update(
+                Panel(
+                    goal_text,
+                    title="Goal",
+                    border_style="magenta",
+                    title_align="left",
+                    padding=(0, 1),
+                    height=3,
+                )
+            )
+
+        step_display = Text()
+
+        if completed:
+            if success:
+                step_display.append("✓ ", style="bold green")
+                panel_title = "Completed"
+                panel_style = "green"
+            else:
+                step_display.append("✗ ", style="bold red")
+                panel_title = "Failed"
+                panel_style = "red"
+        else:
+            step_display.append("⚡ ", style="bold yellow")
+            panel_title = "Status"
+            panel_style = "yellow"
+
+        step_display.append(step_message)
+
+        layout["status"].update(
+            Panel(
+                step_display,
+                title=panel_title,
+                border_style=panel_style,
+                title_align="left",
+                padding=(0, 1),
+                height=3,
+            )
+        )
+
+    def handle_event(self, event):
+        """Handle streaming events from the agent workflow."""
+        logger = logging.getLogger("droidrun")
+
+        # Log different event types with proper names
+        if isinstance(event, ScreenshotEvent):
+            logger.debug("📸 Taking screenshot...")
+
+        # Planner events
+        elif isinstance(event, PlanInputEvent):
+            self.current_step = "Planning..."
+            logger.info("💭 Planner receiving input...")
+
+        elif isinstance(event, PlanThinkingEvent):
+            if event.thoughts:
+                thoughts_preview = (
+                    event.thoughts[:150] + "..."
+                    if len(event.thoughts) > 150
+                    else event.thoughts
+                )
+                logger.info(f"🧠 Planning: {thoughts_preview}")
+            if event.code:
+                logger.info(f"📝 Generated plan code")
+
+        elif isinstance(event, PlanCreatedEvent):
+            if event.tasks:
+                task_count = len(event.tasks) if event.tasks else 0
+                self.current_step = f"Plan ready ({task_count} tasks)"
+                logger.info(f"📋 Plan created with {task_count} tasks")
+                for task in event.tasks:
+                    desc = task.description
+                    logger.info(f"- {desc}")
+
+        # CodeAct events
+        elif isinstance(event, TaskInputEvent):
+            self.current_step = "Processing task input..."
+            logger.info("💬 Task input received...")
+
+        elif isinstance(event, TaskThinkingEvent):
+            if hasattr(event, "thoughts") and event.thoughts:
+                thoughts_preview = (
+                    event.thoughts[:150] + "..."
+                    if len(event.thoughts) > 150
+                    else event.thoughts
+                )
+                logger.info(f"🧠 Thinking: {thoughts_preview}")
+            if hasattr(event, "code") and event.code:
+                logger.info(f"💻 Executing action code")
+                logger.debug(f"{event.code}")
+
+        elif isinstance(event, TaskExecutionEvent):
+            self.current_step = "Executing action..."
+            logger.info(f"⚡ Executing action...")
+
+        elif isinstance(event, TaskExecutionResultEvent):
+            if hasattr(event, "output") and event.output:
+                output = str(event.output)
+                if "Error" in output or "Exception" in output:
+                    output_preview = (
+                        output[:100] + "..." if len(output) > 100 else output
+                    )
+                    logger.info(f"❌ Action error: {output_preview}")
+                else:
+                    output_preview = (
+                        output[:100] + "..." if len(output) > 100 else output
+                    )
+                    logger.info(f"⚡ Action result: {output_preview}")
+
+        elif isinstance(event, TaskEndEvent):
+            if hasattr(event, "success") and hasattr(event, "reason"):
+                if event.success:
+                    self.current_step = event.reason
+                    logger.info(f"✅ Task completed: {event.reason}")
+                else:
+                    self.current_step = f"Task failed"
+                    logger.info(f"❌ Task failed: {event.reason}")
+
+        # Droid coordination events
+        elif isinstance(event, CodeActExecuteEvent):
+            self.current_step = "Executing task..."
+            logger.info(f"🔧 Starting task execution...")
+
+        elif isinstance(event, CodeActResultEvent):
+            if hasattr(event, "success") and hasattr(event, "reason"):
+                if event.success:
+                    self.current_step = event.reason
+                    logger.info(f"✅ Task completed: {event.reason}")
+                else:
+                    self.current_step = f"Task failed"
+                    logger.info(f"❌ Task failed: {event.reason}")
+
+        elif isinstance(event, ReasoningLogicEvent):
+            self.current_step = "Planning..."
+            logger.info(f"🤔 Planning next steps...")
+
+        elif isinstance(event, TaskRunnerEvent):
+            self.current_step = "Processing tasks..."
+            logger.info(f"🏃 Processing task queue...")
+
+        elif isinstance(event, FinalizeEvent):
+            if hasattr(event, "success") and hasattr(event, "reason"):
+                self.is_completed = True
+                self.is_success = event.success
+                if event.success:
+                    self.current_step = f"Success: {event.reason}"
+                    logger.info(f"🎉 Goal achieved: {event.reason}")
+                else:
+                    self.current_step = f"Failed: {event.reason}"
+                    logger.info(f"❌ Goal failed: {event.reason}")
+
+        else:
+            logger.debug(f"🔄 {event.__class__.__name__}")