droidrun 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

droidrun/agent/utils/executer.py

@@ -2,9 +2,16 @@ import io
 import contextlib
 import ast
 import traceback
+import logging
 from typing import Any, Dict
-from .async_utils import async_to_sync
+from droidrun.agent.utils.async_utils import async_to_sync
+from llama_index.core.workflow import Context
 import asyncio
+from asyncio import AbstractEventLoop
+import threading
+
+logger = logging.getLogger("droidrun")
+
 
 class SimpleCodeExecutor:
     """
@@ -16,7 +23,14 @@ class SimpleCodeExecutor:
     NOTE: not safe for production use! Use with caution.
     """
 
-    def __init__(self, loop, locals: Dict[str, Any] = {}, globals: Dict[str, Any] = {}, tools = {}, use_same_scope: bool = True):
+    def __init__(
+        self,
+        loop: AbstractEventLoop,
+        locals: Dict[str, Any] = {},
+        globals: Dict[str, Any] = {},
+        tools={},
+        use_same_scope: bool = True,
+    ):
         """
         Initialize the code executor.
 
@@ -28,9 +42,12 @@ class SimpleCodeExecutor:
 
         # loop throught tools and add them to globals, but before that check if tool value is async, if so convert it to sync. tools is a dictionary of tool name: function
         # e.g. tools = {'tool_name': tool_function}
-        
+
         # check if tools is a dictionary
         if isinstance(tools, dict):
+            logger.debug(
+                f"🔧 Initializing SimpleCodeExecutor with tools: {tools.items()}"
+            )
             for tool_name, tool_function in tools.items():
                 if asyncio.iscoroutinefunction(tool_function):
                     # If the function is async, convert it to sync
@@ -38,6 +55,7 @@ class SimpleCodeExecutor:
                 # Add the tool to globals
                 globals[tool_name] = tool_function
         elif isinstance(tools, list):
+            logger.debug(f"🔧 Initializing SimpleCodeExecutor with tools: {tools}")
             # If tools is a list, convert it to a dictionary with tool name as key and function as value
             for tool in tools:
                 if asyncio.iscoroutinefunction(tool):
@@ -48,20 +66,22 @@ class SimpleCodeExecutor:
         else:
             raise ValueError("Tools must be a dictionary or a list of functions.")
 
-
-        # add time to globals
         import time
-        globals['time'] = time
-        # State that persists between executions
+
+        globals["time"] = time
+
         self.globals = globals
         self.locals = locals
         self.loop = loop
         self.use_same_scope = use_same_scope
         if self.use_same_scope:
             # If using the same scope, set the globals and locals to the same dictionary
-            self.globals = self.locals = {**self.locals, **{k: v for k, v in self.globals.items() if k not in self.locals}}
+            self.globals = self.locals = {
+                **self.locals,
+                **{k: v for k, v in self.globals.items() if k not in self.locals},
+            }
 
-    async def execute(self, code: str) -> str:
+    async def execute(self, ctx: Context, code: str) -> str:
         """
         Execute Python code and capture output and return values.
 
@@ -71,6 +91,9 @@ class SimpleCodeExecutor:
         Returns:
             str: Output from the execution, including print statements.
         """
+        # Update UI elements before execution
+        self.globals['ui_state'] = await ctx.get("ui_state", None)
+        
         # Capture stdout and stderr
         stdout = io.StringIO()
         stderr = io.StringIO()
@@ -78,20 +101,32 @@ class SimpleCodeExecutor:
         output = ""
         try:
             # Execute with captured output
-            with contextlib.redirect_stdout(
-                stdout
-            ), contextlib.redirect_stderr(stderr):
+            thread_exception = []
+            with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
+
+                def execute_code():
+                    try:
+                        exec(code, self.globals, self.locals)
+                    except Exception as e:
+                        import traceback
+
+                        thread_exception.append((e, traceback.format_exc()))
 
-                exec(code, self.globals, self.locals)
+                t = threading.Thread(target=execute_code)
+                t.start()
+                t.join()
 
             # Get output
             output = stdout.getvalue()
             if stderr.getvalue():
                 output += "\n" + stderr.getvalue()
+            if thread_exception:
+                e, tb = thread_exception[0]
+                output += f"\nError: {type(e).__name__}: {str(e)}\n{tb}"
 
         except Exception as e:
             # Capture exception information
             output = f"Error: {type(e).__name__}: {str(e)}\n"
             output += traceback.format_exc()
 
-        return output
+        return output

droidrun/agent/utils/llm_picker.py

@@ -3,8 +3,7 @@ import logging
 from typing import Any
 from llama_index.core.llms.llm import LLM
 # Configure logging
-logger = logging.getLogger(__name__)
-logger.addHandler(logging.NullHandler())
+logger = logging.getLogger("droidrun")
 
 def load_llm(provider_name: str, **kwargs: Any) -> LLM:
     """
@@ -32,6 +31,8 @@ def load_llm(provider_name: str, **kwargs: Any) -> LLM:
         raise ValueError("provider_name cannot be empty.")
     if provider_name == "OpenAILike":
         module_provider_part = "openai_like"
+    elif provider_name == "GoogleGenAI":
+        module_provider_part = "google_genai"
     else:
         # Use lowercase for module path, handle hyphens for package name suggestion
         lower_provider_name = provider_name.lower()
@@ -44,9 +45,9 @@ def load_llm(provider_name: str, **kwargs: Any) -> LLM:
     install_package_name = f"llama-index-llms-{module_provider_part.replace('_', '-')}"
 
     try:
-        logger.info(f"Attempting to import module: {module_path}")
+        logger.debug(f"Attempting to import module: {module_path}")
         llm_module = importlib.import_module(module_path)
-        logger.info(f"Successfully imported module: {module_path}")
+        logger.debug(f"Successfully imported module: {module_path}")
 
     except ModuleNotFoundError:
         logger.error(f"Module '{module_path}' not found. Try: pip install {install_package_name}")
@@ -55,18 +56,21 @@ def load_llm(provider_name: str, **kwargs: Any) -> LLM:
         ) from None
 
     try:
-        logger.info(f"Attempting to get class '{provider_name}' from module {module_path}")
+        logger.debug(f"Attempting to get class '{provider_name}' from module {module_path}")
         llm_class = getattr(llm_module, provider_name)
-        logger.info(f"Found class: {llm_class.__name__}")
+        logger.debug(f"Found class: {llm_class.__name__}")
 
         # Verify the class is a subclass of LLM
         if not isinstance(llm_class, type) or not issubclass(llm_class, LLM):
             raise TypeError(f"Class '{provider_name}' found in '{module_path}' is not a valid LLM subclass.")
 
+        # Filter out None values from kwargs
+        filtered_kwargs = {k: v for k, v in kwargs.items() if v is not None}
+        
         # Initialize
-        logger.info(f"Initializing {llm_class.__name__} with kwargs: {list(kwargs.keys())}")
-        llm_instance = llm_class(**kwargs)
-        logger.info(f"Successfully loaded and initialized LLM: {provider_name}")
+        logger.debug(f"Initializing {llm_class.__name__} with kwargs: {list(filtered_kwargs.keys())}")
+        llm_instance = llm_class(**filtered_kwargs)
+        logger.debug(f"Successfully loaded and initialized LLM: {provider_name}")
         if not llm_instance:
             raise RuntimeError(f"Failed to initialize LLM instance for {provider_name}.")
         return llm_instance
@@ -81,7 +85,7 @@ def load_llm(provider_name: str, **kwargs: Any) -> LLM:
         raise # Re-raise TypeError (could be from issubclass check or __init__)
     except Exception as e:
         logger.error(f"An unexpected error occurred initializing {provider_name}: {e}")
-        raise RuntimeError(f"Failed to initialize LLM '{provider_name}'.") from e
+        raise e
     
 # --- Example Usage ---
 if __name__ == "__main__":

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__}")