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

droidrun/agent/droid/droid_agent.py

@@ -3,40 +3,70 @@ DroidAgent - A wrapper class that coordinates the planning and execution of task
 to achieve a user's goal on an Android device.
 """
 
-import asyncio
 import logging
-from typing import Dict, Any, List, Tuple
+from typing import List
 
-from llama_index.core.base.llms.types import ChatMessage
 from llama_index.core.llms.llm import LLM
-from llama_index.core.memory import ChatMemoryBuffer
-from ..codeact import CodeActAgent
-from ..planner import PlannerAgent, TaskManager
-from ..utils.executer import SimpleCodeExecutor
-from ...tools import Tools
+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:
+class DroidAgent(Workflow):
     """
-    A wrapper class that coordinates between PlannerAgent (creates plans) and 
+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_instance: 'Tools' = None,
-        tool_list: Dict[str, Any] = None,
+        tools: Tools,
+        personas: List[AgentPersona] = [DEFAULT],
         max_steps: int = 15,
-        vision: bool = False,
         timeout: int = 1000,
-        max_retries: int = 3,
-        reasoning: bool = True,
+        reasoning: bool = False,
+        reflection: bool = False,
         enable_tracing: bool = False,
         debug: bool = False,
-        device_serial: str = None,
+        save_trajectories: bool = False,
+        *args,
         **kwargs
     ):
         """
@@ -45,184 +75,83 @@ class DroidAgent:
         Args:
             goal: The user's goal or command to execute
             llm: The language model to use for both agents
-            tools_instance: An instance of the Tools class
-            tool_list: Dictionary of available tools
             max_steps: Maximum number of steps for both agents
-            vision: Whether to enable vision capabilities
             timeout: Timeout for agent execution in seconds
-            max_retries: Maximum number of retries for failed tasks
             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
-            device_serial: Target Android device serial number
             **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.vision = vision
+        self.max_codeact_steps = max_steps
         self.timeout = timeout
-        self.max_retries = max_retries
-        self.task_manager = TaskManager()
         self.reasoning = reasoning
+        self.reflection = reflection
         self.debug = debug
-        self.device_serial = device_serial
-        
-        logger.info("🤖 Initializing DroidAgent wrapper...")
-        
-        self.tools_instance = tools_instance
-        self.tool_list = tool_list
-        
-        # Ensure remember tool is in the tool_list if available
-        if hasattr(tools_instance, 'remember') and 'remember' not in tool_list:
-            logger.debug("📝 Adding 'remember' tool to the available tools")
-            tool_list['remember'] = tools_instance.remember
-        
-        # Create code executor
-        logger.debug("🔧 Initializing Code Executor...")
-        loop = asyncio.get_event_loop()
-        self.executor = SimpleCodeExecutor(
-            loop=loop,
-            locals={},
-            tools=tool_list,
-            globals={"__builtins__": __builtins__}
-        )
-        
-        # Create memory buffer for the planning agent if reasoning is enabled
-        if self.reasoning:
-            self.planning_memory = ChatMemoryBuffer.from_defaults(llm=llm)
+
+        self.event_counter = 0
+        self.save_trajectories = save_trajectories
         
-        # Create CodeActAgent
-        logger.info("🧠 Initializing CodeAct Agent...")
-        self.codeact_agent = CodeActAgent(
-            llm=llm,
-            code_execute_fn=self.executor.execute,
-            available_tools=tool_list.values(),
-            tools=tools_instance,
-            max_steps=999999, 
-            vision=vision,
-            debug=debug,
-            timeout=timeout
-        )
+        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,
-                agent=self.codeact_agent, 
-                tools_instance=tools_instance,
+                personas=personas,
+                task_manager=self.task_manager,
+                tools_instance=tools,
                 timeout=timeout,
-                max_retries=max_retries,
-                enable_tracing=enable_tracing,
                 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)
             
-            # Give task manager to the planner
-            self.planner_agent.task_manager = self.task_manager
         else:
             logger.debug("🚫 Planning disabled - will execute tasks directly with CodeActAgent")
             self.planner_agent = None
         
         logger.info("✅ DroidAgent initialized successfully.")
     
-    async def _get_plan_from_planner(self) -> List[Dict]:
-        """
-        Get a plan (list of tasks) from the PlannerAgent.
-        
-        Returns:
-            List of task dictionaries
-        """
-        logger.info("📋 Planning steps to accomplish the goal...")
-        
-        # Create system and user messages
-        system_msg = ChatMessage(role="system", content=self.planner_agent.system_prompt)
-        user_msg = ChatMessage(role="user", content=self.planner_agent.user_prompt)
-        
-        # Check if we have task history to add to the prompt
-        task_history = ""
-        # Use the persistent task history methods to get ALL completed and failed tasks
-        completed_tasks = self.task_manager.get_all_completed_tasks()
-        failed_tasks = self.task_manager.get_all_failed_tasks()
-        
-        # Show any remembered information in task history
-        remembered_info = ""
-        if hasattr(self.tools_instance, 'memory') and self.tools_instance.memory:
-            remembered_info = "\n### Remembered Information:\n"
-            for idx, item in enumerate(self.tools_instance.memory, 1):
-                remembered_info += f"{idx}. {item}\n"
-        
-        if completed_tasks or failed_tasks or remembered_info:
-            task_history = "### Task Execution History:\n"
-            
-            if completed_tasks:
-                task_history += "✅ Completed Tasks:\n"
-                for task in completed_tasks:
-                    task_history += f"- {task['description']}\n"
-
-            if failed_tasks:
-                task_history += "\n❌ Failed Tasks:\n"
-                for task in failed_tasks:
-                    failure_reason = task.get('failure_reason', 'Unknown reason')
-                    task_history += f"- {task['description']} (Failed: {failure_reason})\n"
-            
-            if remembered_info:
-                task_history += remembered_info
-                
-            # Add a reminder to use this information
-            task_history += "\n⚠️ Please use the above information in your planning. For example, if specific dates or locations were found, include them explicitly in your next tasks instead of just referring to 'the dates' or 'the location'.\n"
-            
-            # Append task history to user prompt
-            user_msg = ChatMessage(
-                role="user", 
-                content=f"{self.planner_agent.user_prompt}\n\n{task_history}\n\nPlease consider the above task history and discovered information when creating your next plan. Incorporate specific data (dates, locations, etc.) directly into tasks rather than referring to them generally. Remember that previously completed or failed tasks will not be repeated."
-            )
-        
-        # Create message list
-        messages = [system_msg, user_msg]
-        if self.debug:
-            logger.debug(f"Sending {len(messages)} messages to planner: {[msg.role for msg in messages]}")
-        
-        # Get response from LLM
-        llm_response = await self.planner_agent._get_llm_response(messages)
-        code, thoughts = self.planner_agent._extract_code_and_thought(llm_response.message.content)
-        
-        # Execute the planning code (which should call set_tasks)
-        if code:
-            try:
-                planning_tools = {
-                    "set_tasks": self.task_manager.set_tasks,
-                    "add_task": self.task_manager.add_task,
-                    "get_all_tasks": self.task_manager.get_all_tasks,
-                    "clear_tasks": self.task_manager.clear_tasks,
-                    "complete_goal": self.task_manager.complete_goal
-                }
-                planning_executor = SimpleCodeExecutor(
-                    loop=asyncio.get_event_loop(),
-                    globals={},
-                    locals={},
-                    tools=planning_tools
-                )
-                await planning_executor.execute(code)
-            except Exception as e:
-                logger.error(f"Error executing planning code: {e}")
-                # If there's an error, create a simple default task
-                self.task_manager.set_tasks([f"Achieve the goal: {self.goal}"])
-        
-        # Get and display the tasks
-        tasks = self.task_manager.get_all_tasks()
-        if tasks:
-            logger.info("📝 Plan created:")
-            for i, task in enumerate(tasks, 1):
-                if task["status"] == self.task_manager.STATUS_PENDING:
-                    logger.info(f"  {i}. {task['description']}")
-        else:
-            logger.warning("No tasks were generated in the plan")
-            
-        return tasks
-
-    async def _execute_task_with_codeact(self, task: Dict) -> Tuple[bool, str]:
+    @step
+    async def execute_task(
+        self,
+        ctx: Context,
+        ev: CodeActExecuteEvent
+        ) -> CodeActResultEvent:
         """
         Execute a single task using the CodeActAgent.
         
@@ -232,59 +161,144 @@ class DroidAgent:
         Returns:
             Tuple of (success, reason)
         """
-        task_description = task["description"]
-        logger.info(f"🔧 Executing task: {task_description}")
-        
-        # Update task status
-        task["status"] = self.task_manager.STATUS_ATTEMPTING
-        
-        # Run the CodeActAgent
+        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:
-            # Reset the tools finished flag before execution
-            self.tools_instance.finished = False
-            self.tools_instance.success = None
-            self.tools_instance.reason = None
-            
-            # Execute the CodeActAgent with the task description as input
-            # Pass input as a keyword argument, not as a dictionary
-            result = await self.codeact_agent.run(input=task_description)
+            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
+            )
             
-            # Check if the tools instance was marked as finished by the 'complete' function
-            if self.tools_instance.finished:
-                if self.tools_instance.success:
-                    task["status"] = self.task_manager.STATUS_COMPLETED
-                    if self.debug:
-                        logger.debug(f"Task completed successfully: {self.tools_instance.reason}")
-                    return True, self.tools_instance.reason or "Task completed successfully"
-                else:
-                    task["status"] = self.task_manager.STATUS_FAILED
-                    task["failure_reason"] = self.tools_instance.reason or "Task failed without specific reason"
-                    logger.warning(f"Task failed: {task['failure_reason']}")
-                    return False, self.tools_instance.reason or "Task failed without specific reason"
+            async for nested_ev in handler.stream_events():
+                self.handle_stream_event(nested_ev, ctx)
+
+            result = await handler
+
             
-            # If tools instance wasn't marked as finished, check the result directly
-            if result and isinstance(result, dict) and "success" in result and result["success"]:
-                task["status"] = self.task_manager.STATUS_COMPLETED
-                if self.debug:
-                    logger.debug(f"Task completed with result: {result}")
-                return True, result.get("reason", "Task completed successfully")
+            if "success" in result and result["success"]:
+                return CodeActResultEvent(success=True, reason=result["reason"], task=task, steps=result["codeact_steps"])
             else:
-                failure_reason = result.get("reason", "Unknown failure") if isinstance(result, dict) else "Task execution failed"
-                task["status"] = self.task_manager.STATUS_FAILED
-                task["failure_reason"] = failure_reason
-                logger.warning(f"Task failed: {failure_reason}")
-                return False, failure_reason
+                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())
-            task["status"] = self.task_manager.STATUS_FAILED
-            task["failure_reason"] = f"Error: {str(e)}"
-            return False, f"Error: {str(e)}"
+            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()
 
-    async def run(self) -> Dict[str, Any]:
+        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.
         
@@ -293,126 +307,51 @@ class DroidAgent:
         """
         logger.info(f"🚀 Running DroidAgent to achieve goal: {self.goal}")
         
-        step_counter = 0
-        retry_counter = 0
-        overall_success = False
-        final_message = ""
-        
-        try:
-            # If reasoning is disabled, directly execute the goal as a single task in CodeActAgent
-            if not self.reasoning:
-                logger.info(f"🔄 Direct execution mode - executing goal: {self.goal}")
-                # Create a simple task for the goal
-                task = {
-                    "description": self.goal,
-                    "status": self.task_manager.STATUS_PENDING
-                }
-                
-                # Execute the task directly with CodeActAgent
-                success, reason = await self._execute_task_with_codeact(task)
-                
-                return {
-                    "success": success,
-                    "reason": reason,
-                    "steps": 1,
-                    "task_history": [task]  # Single task history
-                }
+        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"
+            )
             
-            # Standard reasoning mode with planning
-            while step_counter < self.max_steps:
-                step_counter += 1
-                if self.debug:
-                    logger.debug(f"Planning step {step_counter}/{self.max_steps}")
-                
-                # 1. Get a plan from the planner
-                tasks = await self._get_plan_from_planner()
-                
-                if self.task_manager.task_completed:
-                    # Task is marked as complete by the planner
-                    logger.info(f"✅ Goal completed: {self.task_manager.message}")
-                    overall_success = True
-                    final_message = self.task_manager.message
-                    break
-                
-                if not tasks:
-                    logger.warning("No tasks generated by planner")
-                    final_message = "Planner did not generate any tasks"
-                    break
-                
-                # 2. Execute each task in the plan sequentially
-                for task in tasks:
-                    if task["status"] == self.task_manager.STATUS_PENDING:
-                        # Reset the CodeActAgent's step counter for this task
-                        self.codeact_agent.steps_counter = 0
-                        
-                        # Execute the task
-                        success, reason = await self._execute_task_with_codeact(task)
-                        
-                        # Update task info with detailed result for the planner
-                        task_idx = tasks.index(task)
-                        result_info = {
-                            "execution_details": reason,
-                            "step_executed": step_counter,
-                            "codeact_steps": self.codeact_agent.steps_counter
-                        }
-                        
-                        # Only update if not already updated in _execute_task_with_codeact
-                        if success:
-                            self.task_manager.update_status(
-                                task_idx, 
-                                self.task_manager.STATUS_COMPLETED, 
-                                result_info
-                            )
-                            logger.info(f"✅ Task completed: {task['description']}")
-                        
-                        if not success:
-                            # Store detailed failure information if not already set
-                            if "failure_reason" not in task:
-                                self.task_manager.update_status(
-                                    task_idx,
-                                    self.task_manager.STATUS_FAILED,
-                                    {"failure_reason": reason, **result_info}
-                                )
-                            
-                            # Handle retries
-                            if retry_counter < self.max_retries:
-                                retry_counter += 1
-                                logger.info(f"Retrying... ({retry_counter}/{self.max_retries})")
-                                # Next iteration will generate a new plan based on current state
-                                break
-                            else:
-                                logger.error(f"Max retries exceeded for task")
-                                final_message = f"Failed after {self.max_retries} retries. Reason: {reason}"
-                                return {"success": False, "reason": final_message}
-                
-                # Reset retry counter for new task sequence
-                retry_counter = 0
+            return CodeActExecuteEvent(task=task, reflection=None)
+        
+        return ReasoningLogicEvent()
                 
-                # Check if all tasks are completed
-                all_completed = all(task["status"] == self.task_manager.STATUS_COMPLETED for task in tasks)
-                if all_completed:
-                    # Get a new plan (the planner might decide we're done)
-                    continue
+        
+    @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)
             
-            # Check if we exited due to max steps
-            if step_counter >= self.max_steps and not overall_success:
-                final_message = f"Reached maximum number of steps ({self.max_steps})"
-                overall_success = False
-                
-            return {
-                "success": overall_success,
-                "reason": final_message,
-                "steps": step_counter,
-                "task_history": self.task_manager.get_task_history()
-            }
-                
-        except Exception as e:
-            logger.error(f"❌ Error during DroidAgent execution: {e}")
-            if self.debug:
-                import traceback
-                logger.error(traceback.format_exc())
-            return {
-                "success": False, 
-                "reason": str(e),
-                "task_history": self.task_manager.get_task_history()
-            } 
+            if isinstance(ev, ScreenshotEvent):
+                self.trajectory.screenshots.append(ev.screenshot)
+
+            else:
+                self.trajectory.events.append(ev)
+
+
+