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

droidrun/agent/droid/droid_agent.py

@@ -0,0 +1,357 @@
+"""
+DroidAgent - A wrapper class that coordinates the planning and execution of tasks
+to achieve a user's goal on an Android device.
+"""
+
+import logging
+from typing import List
+
+from llama_index.core.llms.llm import LLM
+from llama_index.core.workflow import step, StartEvent, StopEvent, Workflow, Context
+from droidrun.agent.droid.events import *
+from droidrun.agent.codeact import CodeActAgent
+from droidrun.agent.codeact.events import EpisodicMemoryEvent
+from droidrun.agent.planner import PlannerAgent
+from droidrun.agent.context.task_manager import TaskManager
+from droidrun.agent.utils.trajectory import Trajectory
+from droidrun.tools import Tools, describe_tools
+from droidrun.agent.common.events import ScreenshotEvent
+from droidrun.agent.common.default import MockWorkflow
+from droidrun.agent.context import ContextInjectionManager
+from droidrun.agent.context.agent_persona import AgentPersona
+from droidrun.agent.context.personas import DEFAULT
+from droidrun.agent.oneflows.reflector import Reflector
+
+
+logger = logging.getLogger("droidrun")
+
+class DroidAgent(Workflow):
+    """
+A wrapper class that coordinates between PlannerAgent (creates plans) and 
+    CodeActAgent (executes tasks) to achieve a user's goal.
+    """
+    
+    @staticmethod
+    def _configure_default_logging(debug: bool = False):
+        """
+        Configure default logging for DroidAgent if no handlers are present.
+        This ensures logs are visible when using DroidAgent directly.
+        """
+        # Only configure if no handlers exist (avoid duplicate configuration)
+        if not logger.handlers:
+            # Create a console handler
+            handler = logging.StreamHandler()
+            
+            # Set format
+            if debug:
+                formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s", "%H:%M:%S")
+            else:
+                formatter = logging.Formatter("%(message)s")
+            
+            handler.setFormatter(formatter)
+            logger.addHandler(handler)
+            logger.setLevel(logging.DEBUG if debug else logging.INFO)
+            logger.propagate = False
+    
+    def __init__(
+        self, 
+        goal: str,
+        llm: LLM,
+        tools: Tools,
+        personas: List[AgentPersona] = [DEFAULT],
+        max_steps: int = 15,
+        timeout: int = 1000,
+        reasoning: bool = False,
+        reflection: bool = False,
+        enable_tracing: bool = False,
+        debug: bool = False,
+        save_trajectories: bool = False,
+        *args,
+        **kwargs
+    ):
+        """
+        Initialize the DroidAgent wrapper.
+        
+        Args:
+            goal: The user's goal or command to execute
+            llm: The language model to use for both agents
+            max_steps: Maximum number of steps for both agents
+            timeout: Timeout for agent execution in seconds
+            reasoning: Whether to use the PlannerAgent for complex reasoning (True) 
+                      or send tasks directly to CodeActAgent (False)
+            reflection: Whether to reflect on steps the CodeActAgent did to give the PlannerAgent advice
+            enable_tracing: Whether to enable Arize Phoenix tracing
+            debug: Whether to enable verbose debug logging
+            **kwargs: Additional keyword arguments to pass to the agents
+        """
+        super().__init__(timeout=timeout ,*args,**kwargs)
+        
+        # Configure default logging if not already configured
+        self._configure_default_logging(debug=debug)
+        
+        # Setup global tracing first if enabled
+        if enable_tracing:
+            try:
+                from llama_index.core import set_global_handler
+                set_global_handler("arize_phoenix")
+                logger.info("🔍 Arize Phoenix tracing enabled globally")
+            except ImportError:
+                logger.warning("⚠️ Arize Phoenix package not found, tracing disabled")
+                enable_tracing = False
+
+        self.goal = goal
+        self.llm = llm
+        self.max_steps = max_steps
+        self.max_codeact_steps = max_steps
+        self.timeout = timeout
+        self.reasoning = reasoning
+        self.reflection = reflection
+        self.debug = debug
+
+        self.event_counter = 0
+        self.save_trajectories = save_trajectories
+        
+        self.trajectory = Trajectory()
+        self.task_manager = TaskManager()
+        self.task_iter = None
+        self.cim = ContextInjectionManager(personas=personas)
+        self.current_episodic_memory = None
+
+        logger.info("🤖 Initializing DroidAgent...")
+        
+        self.tool_list = describe_tools(tools)
+        self.tools_instance = tools
+
+
+        if self.reasoning:
+            logger.info("📝 Initializing Planner Agent...")
+            self.planner_agent = PlannerAgent(
+                goal=goal,
+                llm=llm,
+                personas=personas,
+                task_manager=self.task_manager,
+                tools_instance=tools,
+                timeout=timeout,
+                debug=debug
+            )
+            self.add_workflows(planner_agent=self.planner_agent)
+            self.max_codeact_steps = 5
+
+            if self.reflection:
+                self.reflector = Reflector(llm=llm, debug=debug)
+            
+        else:
+            logger.debug("🚫 Planning disabled - will execute tasks directly with CodeActAgent")
+            self.planner_agent = None
+        
+        logger.info("✅ DroidAgent initialized successfully.")
+    
+    @step
+    async def execute_task(
+        self,
+        ctx: Context,
+        ev: CodeActExecuteEvent
+        ) -> CodeActResultEvent:
+        """
+        Execute a single task using the CodeActAgent.
+        
+        Args:
+            task: Task dictionary with description and status
+            
+        Returns:
+            Tuple of (success, reason)
+        """
+        task: Task = ev.task
+        reflection = ev.reflection if ev.reflection is not None else None
+        persona = self.cim.get_persona(task.agent_type)
+
+        logger.info(f"🔧 Executing task: {task.description}")
+
+        try:
+            codeact_agent = CodeActAgent(
+                llm=self.llm,
+                persona=persona,
+                max_steps=self.max_codeact_steps,
+                all_tools_list=self.tool_list,
+                tools_instance=self.tools_instance,
+                debug=self.debug,
+                timeout=self.timeout,
+            )
+
+            handler = codeact_agent.run(
+                input=task.description,
+                remembered_info=self.tools_instance.memory,
+                reflection=reflection
+            )
+            
+            async for nested_ev in handler.stream_events():
+                self.handle_stream_event(nested_ev, ctx)
+
+            result = await handler
+
+            
+            if "success" in result and result["success"]:
+                return CodeActResultEvent(success=True, reason=result["reason"], task=task, steps=result["codeact_steps"])
+            else:
+                return CodeActResultEvent(success=False, reason=result["reason"], task=task, steps=result["codeact_steps"])
+                
+        except Exception as e:
+            logger.error(f"Error during task execution: {e}")
+            if self.debug:
+                import traceback
+                logger.error(traceback.format_exc())
+            return CodeActResultEvent(success=False, reason=f"Error: {str(e)}", task=task, steps=result["codeact_steps"])
+    
+    @step
+    async def handle_codeact_execute(self, ctx: Context, ev: CodeActResultEvent) -> FinalizeEvent | ReflectionEvent:
+        try:
+            task = ev.task
+            if not self.reasoning:
+                return FinalizeEvent(success=ev.success, reason=ev.reason, task=[task], steps=ev.steps)
+            
+            if self.reflection:
+                return ReflectionEvent(task=task)
+            
+            return ReasoningLogicEvent()
+
+        except Exception as e:
+            logger.error(f"❌ Error during DroidAgent execution: {e}")
+            if self.debug:
+                import traceback
+                logger.error(traceback.format_exc())
+            return FinalizeEvent(success=False, reason=str(e), task=self.task_manager.get_task_history(), steps=self.step_counter)
+        
+
+    @step
+    async def reflect(
+        self,
+        ctx: Context,
+        ev: ReflectionEvent
+        ) -> ReasoningLogicEvent | CodeActExecuteEvent:
+
+
+        task = ev.task
+        if ev.task.agent_type == "AppStarterExpert":
+            self.task_manager.complete_task(task)
+            return ReasoningLogicEvent()
+        
+        reflection = await self.reflector.reflect_on_episodic_memory(episodic_memory=self.current_episodic_memory, goal=task.description)
+
+        if reflection.goal_achieved:
+            self.task_manager.complete_task(task)
+            return ReasoningLogicEvent()
+        
+        else:
+            self.task_manager.fail_task(task)
+            return ReasoningLogicEvent(reflection=reflection)
+        
+
+    @step
+    async def handle_reasoning_logic(
+        self,
+        ctx: Context,
+        ev: ReasoningLogicEvent,
+        planner_agent: Workflow = MockWorkflow()
+        ) -> FinalizeEvent | CodeActExecuteEvent:
+        try:
+            if self.step_counter >= self.max_steps:
+                return FinalizeEvent(success=False, reason=f"Reached maximum number of steps ({self.max_steps})", task=self.task_manager.get_task_history(), steps=self.step_counter)
+            self.step_counter += 1
+
+            if ev.reflection:
+                handler = planner_agent.run(remembered_info=self.tools_instance.memory, reflection=ev.reflection)
+            else:
+                if self.task_iter:
+                    try:
+                        task = next(self.task_iter)
+                        return CodeActExecuteEvent(task=task, reflection=None)
+                    except StopIteration as e:
+                        logger.info("Planning next steps...")
+
+                logger.debug(f"Planning step {self.step_counter}/{self.max_steps}")
+
+                handler = planner_agent.run(remembered_info=self.tools_instance.memory, reflection=None)
+
+            async for nested_ev in handler.stream_events():
+                self.handle_stream_event(nested_ev, ctx)
+
+            result = await handler
+
+            self.tasks = self.task_manager.get_all_tasks()
+            self.task_iter = iter(self.tasks)
+
+            if self.task_manager.goal_completed:
+                logger.info(f"✅ Goal completed: {self.task_manager.message}")
+                return FinalizeEvent(success=True, reason=self.task_manager.message, task=self.task_manager.get_task_history(), steps=self.step_counter)
+            if not self.tasks:
+                logger.warning("No tasks generated by planner")
+                return FinalizeEvent(success=False, reason="Planner did not generate any tasks", task=self.task_manager.get_task_history(), steps=self.step_counter)
+            
+            return CodeActExecuteEvent(task=next(self.task_iter), reflection=None)
+        
+        except Exception as e:
+            logger.error(f"❌ Error during DroidAgent execution: {e}")
+            if self.debug:
+                import traceback
+                logger.error(traceback.format_exc())
+            return FinalizeEvent(success=False, reason=str(e), task=self.task_manager.get_task_history(), steps=self.step_counter)
+    
+
+    @step
+    async def start_handler(self, ctx: Context, ev: StartEvent) -> CodeActExecuteEvent | ReasoningLogicEvent:
+        """
+        Main execution loop that coordinates between planning and execution.
+        
+        Returns:
+            Dict containing the execution result
+        """
+        logger.info(f"🚀 Running DroidAgent to achieve goal: {self.goal}")
+        
+        self.step_counter = 0
+        self.retry_counter = 0
+    
+        if not self.reasoning:
+            logger.info(f"🔄 Direct execution mode - executing goal: {self.goal}")
+            task = Task(
+                description=self.goal,
+                status=self.task_manager.STATUS_PENDING,
+                agent_type="Default"
+            )
+            
+            return CodeActExecuteEvent(task=task, reflection=None)
+        
+        return ReasoningLogicEvent()
+                
+        
+    @step
+    async def finalize(self, ctx: Context, ev: FinalizeEvent) -> StopEvent:
+        ctx.write_event_to_stream(ev)
+
+        result = {
+            "success": ev.success,
+            "reason": ev.reason,
+            "steps": ev.steps,
+        }
+
+        if self.trajectory and self.save_trajectories:
+            self.trajectory.save_trajectory()
+
+        return StopEvent(result)
+    
+    def handle_stream_event(self, ev: Event, ctx: Context):
+
+        if isinstance(ev, EpisodicMemoryEvent):
+            self.current_episodic_memory = ev.episodic_memory
+            return
+
+        if not isinstance(ev, StopEvent):
+            ctx.write_event_to_stream(ev)
+            
+            if isinstance(ev, ScreenshotEvent):
+                self.trajectory.screenshots.append(ev.screenshot)
+
+            else:
+                self.trajectory.events.append(ev)
+
+
+

droidrun/agent/droid/events.py

@@ -0,0 +1,28 @@
+from llama_index.core.workflow import Event
+from droidrun.agent.context import Reflection, Task
+from typing import List, Optional
+
+class CodeActExecuteEvent(Event):
+    task: Task
+    reflection: Optional[Reflection]
+
+class CodeActResultEvent(Event):
+    success: bool
+    reason: str
+    steps: int
+
+class ReasoningLogicEvent(Event):
+    reflection: Optional[Reflection] = None
+
+class FinalizeEvent(Event):
+    success: bool
+    reason: str
+    task: List[Task]
+    steps: int = 1
+
+class TaskRunnerEvent(Event):
+    pass
+
+class ReflectionEvent(Event):
+    task: Task
+    pass

droidrun/agent/oneflows/reflector.py

@@ -0,0 +1,265 @@
+from llama_index.core.llms.llm import LLM
+from droidrun.agent.context import EpisodicMemory
+from droidrun.agent.context.reflection import Reflection
+from llama_index.core.base.llms.types import ChatMessage, ImageBlock
+from droidrun.agent.utils.chat_utils import add_screenshot_image_block
+from droidrun.agent.context.agent_persona import AgentPersona
+import json
+from typing import Dict, Any, List, Optional
+import logging
+from PIL import Image, ImageDraw, ImageFont
+import io
+
+logger = logging.getLogger("droidrun")
+
+class Reflector:
+    def __init__(
+        self,
+        llm: LLM,
+        debug: bool = False,
+        *args,
+        **kwargs
+    ):
+        self.llm = llm
+        self.debug = debug
+
+    async def reflect_on_episodic_memory(self, episodic_memory: EpisodicMemory, goal: str) -> Reflection:
+        """Analyze episodic memory and provide reflection on the agent's performance."""
+        system_prompt_content = self._create_system_prompt()
+        system_prompt = ChatMessage(role="system", content=system_prompt_content)
+
+        episodic_memory_content = self._format_episodic_memory(episodic_memory)
+        persona_content = self._format_persona(episodic_memory.persona)
+        
+        # Create user message content with persona information
+        user_content = f"{persona_content}\n\nGoal: {goal}\n\nEpisodic Memory Steps:\n{episodic_memory_content}\n\nPlease evaluate if the goal was achieved and provide your analysis in the specified JSON format."
+        
+        # Create user message
+        user_message = ChatMessage(role="user", content=user_content)
+        
+        # Create the screenshots grid and add as ImageBlock if screenshots exist
+        screenshots_grid = self._create_screenshots_grid(episodic_memory)
+        
+        if screenshots_grid:
+            # Use the add_screenshot_image_block function to properly add the image
+            messages_list = [system_prompt, user_message]
+            messages_list = await add_screenshot_image_block(screenshots_grid, messages_list, copy=False)
+            messages = messages_list
+        else:
+            messages = [system_prompt, user_message]
+        response = await self.llm.achat(messages=messages)
+
+        logger.info(f"REFLECTION {response.message.content}")
+        
+        try:
+            # Clean the response content to handle markdown code blocks
+            content = response.message.content.strip()
+            
+            # Remove markdown code block formatting if present
+            if content.startswith('```json'):
+                content = content[7:]  # Remove ```json
+            elif content.startswith('```'):
+                content = content[3:]   # Remove ```
+            
+            if content.endswith('```'):
+                content = content[:-3]  # Remove trailing ```
+            
+            content = content.strip()
+            
+            parsed_response = json.loads(content)
+            return Reflection.from_dict(parsed_response)
+        except json.JSONDecodeError as e:
+            logger.error(f"Failed to parse reflection response: {e}")
+            logger.error(f"Raw response: {response.message.content}")
+            return await self.reflect_on_episodic_memory(episodic_memory=episodic_memory, goal=goal)
+    
+    def _create_screenshots_grid(self, episodic_memory: EpisodicMemory) -> Optional[bytes]:
+        """Create a 3x2 grid of screenshots from episodic memory steps."""
+        # Extract screenshots from steps
+        screenshots = []
+        for step in episodic_memory.steps:
+            if step.screenshot:
+                try:
+                    # Convert bytes to PIL Image
+                    screenshot_image = Image.open(io.BytesIO(step.screenshot))
+                    screenshots.append(screenshot_image)
+                except Exception as e:
+                    logger.warning(f"Failed to load screenshot: {e}")
+                    continue
+        
+        if not screenshots:
+            return None
+        
+        num_screenshots = min(len(screenshots), 6)
+        cols, rows = num_screenshots, 1
+        
+        screenshots = screenshots[:num_screenshots]
+        
+        if not screenshots:
+            return None
+        
+        if screenshots:
+            cell_width = screenshots[0].width // 2
+            cell_height = screenshots[0].height // 2
+        else:
+            return None
+        
+        # Define header bar height
+        header_height = 60
+        
+        # Create the grid image with space for header bars
+        grid_width = cols * cell_width
+        grid_height = rows * (cell_height + header_height)
+        grid_image = Image.new('RGB', (grid_width, grid_height), color='white')
+        
+        # Set up font for step text
+        draw = ImageDraw.Draw(grid_image)
+        try:
+            # Use larger font for header text
+            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 48)
+        except:
+            font = ImageFont.load_default()
+        
+        # Place screenshots in the grid with header bars
+        for i, screenshot in enumerate(screenshots):
+            row = i // cols
+            col = i % cols
+            
+            # Calculate positions
+            x = col * cell_width
+            header_y = row * (cell_height + header_height)
+            screenshot_y = header_y + header_height
+            
+            # Create header bar
+            header_rect = [x, header_y, x + cell_width, header_y + header_height]
+            draw.rectangle(header_rect, fill='#2c3e50')  # Dark blue header
+            
+            # Draw step text in header bar
+            text = f"Step {i+1}"
+            # Get text dimensions for centering
+            bbox = draw.textbbox((0, 0), text, font=font)
+            text_width = bbox[2] - bbox[0]
+            text_height = bbox[3] - bbox[1]
+            
+            # Center text in header bar
+            text_x = x + (cell_width - text_width) // 2
+            text_y = header_y + (header_height - text_height) // 2
+            
+            draw.text((text_x, text_y), text, fill='white', font=font)
+            
+            # Resize and place screenshot below header
+            resized_screenshot = screenshot.resize((cell_width, cell_height), Image.Resampling.LANCZOS)
+            grid_image.paste(resized_screenshot, (x, screenshot_y))
+        
+        # Save grid to disk for debugging (only if debug flag is enabled)
+        if self.debug:
+            import os
+            from datetime import datetime
+            
+            # Create debug directory if it doesn't exist
+            debug_dir = "reflection_screenshots"
+            os.makedirs(debug_dir, exist_ok=True)
+            
+            # Save with timestamp
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            debug_filename = os.path.join(debug_dir, f"screenshot_grid_{timestamp}.png")
+            grid_image.save(debug_filename)
+            logger.info(f"Screenshot grid saved to: {debug_filename}")
+        
+        # Convert to bytes for use with add_screenshot_image_block
+        buffer = io.BytesIO()
+        grid_image.save(buffer, format='PNG')
+        buffer.seek(0)
+        
+        return buffer.getvalue()
+
+    def _create_system_prompt(self) -> str:
+        """Create a system prompt with reflection instructions."""
+        system_prompt = """You are a Reflector AI that analyzes the performance of an Android Agent. Your role is to examine episodic memory steps and evaluate whether the agent achieved its goal.
+
+        EVALUATION PROCESS:
+        1. First, determine if the agent achieved the stated goal based on the episodic memory steps
+        2. If the goal was achieved, acknowledge the success
+        3. If the goal was NOT achieved, analyze what went wrong and provide direct advice
+        4. Use the provided screenshots (if any) to understand the visual context of each step
+        The screenshots show a screen the agent saw. It is in chronological order from left to right
+
+        ANALYSIS AREAS (for failed goals):
+        - Missed opportunities or inefficient actions
+        - Incorrect tool usage or navigation choices
+        - Failure to understand context or user intent
+        - Suboptimal decision-making patterns
+
+        ADVICE GUIDELINES (for failed goals):
+        - Address the agent directly using "you" form with present/future focus (e.g., "You need to...", "Look for...", "Focus on...")
+        - Provide situational awareness advice that helps with the current state after the failed attempt
+        - Give actionable guidance for what to do NOW when retrying the goal, not what went wrong before
+        - Consider the current app state and context the agent will face when retrying
+        - Focus on the key strategy or approach needed for success in the current situation
+        - Keep it concise but precise (1-2 sentences)
+
+        OUTPUT FORMAT:
+        You MUST respond with a valid JSON object in this exact format:
+
+        {{
+            "goal_achieved": true,
+            "advice": null,
+            "summary": "Brief summary of what happened"
+        }}
+
+        OR
+
+        {{
+            "goal_achieved": false,
+            "advice": "Direct advice using 'you' form focused on current situation - what you need to do NOW when retrying",
+            "summary": "Brief summary of what happened"
+        }}
+
+        IMPORTANT: 
+        - If goal_achieved is true, set advice to null
+        - If goal_achieved is false, provide direct "you" form advice focused on what to do NOW in the current situation when retrying
+        - Advice should be forward-looking and situational, not retrospective about past mistakes
+        - Always include a brief summary of the agent's performance
+        - Ensure the JSON is valid and parsable
+        - ONLY return the JSON object, no additional text or formatting"""
+
+        return system_prompt
+    
+    def _format_persona(self, persona: AgentPersona) -> str:
+        """Format the agent persona information for the user prompt."""
+        persona_content = f"""ACTOR AGENT PERSONA:
+        - Name: {persona.name}
+        - Description: {persona.description}
+        - Available Tools: {', '.join(persona.allowed_tools)}
+        - Expertise Areas: {', '.join(persona.expertise_areas)}
+        - System Prompt: {persona.system_prompt}"""
+                
+        return persona_content
+    
+    def _format_episodic_memory(self, episodic_memory: EpisodicMemory) -> str:
+        """Format the episodic memory steps into a readable format for analysis."""
+        formatted_steps = []
+        
+        for i, step in enumerate(episodic_memory.steps, 1):
+            try:
+                # Parse the JSON strings to get the original content without escape characters
+                chat_history = json.loads(step.chat_history)
+                response = json.loads(step.response)
+                
+                
+                formatted_step = f"""Step {i}:
+            Chat History: {json.dumps(chat_history, indent=2)}
+            Response: {json.dumps(response, indent=2)}
+            Timestamp: {step.timestamp}
+            ---"""
+            except json.JSONDecodeError as e:
+                # Fallback to original format if JSON parsing fails
+                logger.warning(f"Failed to parse JSON for step {i}: {e}")
+                formatted_step = f"""Step {i}:
+            Chat History: {step.chat_history}
+            Response: {step.response}
+            Timestamp: {step.timestamp}
+            ---"""
+            formatted_steps.append(formatted_step)
+        
+        return "\n".join(formatted_steps)

droidrun/agent/planner/__init__.py

@@ -0,0 +1,13 @@
+from droidrun.agent.planner.planner_agent import PlannerAgent
+from droidrun.agent.planner.prompts import (
+    DEFAULT_PLANNER_SYSTEM_PROMPT,
+    DEFAULT_PLANNER_USER_PROMPT,
+    DEFAULT_PLANNER_TASK_FAILED_PROMPT
+)
+
+__all__ = [
+    "PlannerAgent", 
+    "DEFAULT_PLANNER_SYSTEM_PROMPT",
+    "DEFAULT_PLANNER_USER_PROMPT",
+    "DEFAULT_PLANNER_TASK_FAILED_PROMPT"
+]

droidrun/agent/planner/events.py

@@ -0,0 +1,16 @@
+from llama_index.core.workflow import Event
+from llama_index.core.base.llms.types import ChatMessage
+from typing import Optional, Any
+from droidrun.agent.context import Task
+
+class PlanInputEvent(Event):
+    input: list[ChatMessage]
+
+
+class PlanThinkingEvent(Event):
+    thoughts: Optional[str] = None
+    code: Optional[str] = None  
+
+
+class PlanCreatedEvent(Event):
+    tasks: list[Task]