droidrun 0.3.9__py3-none-any.whl → 0.3.10.dev3__py3-none-any.whl

droidrun/agent/droid/droid_agent.py

@@ -1,41 +1,63 @@
 """
 DroidAgent - A wrapper class that coordinates the planning and execution of tasks
 to achieve a user's goal on an Android device.
+
+Architecture:
+- When reasoning=False: Uses CodeActAgent directly
+- When reasoning=True: Uses Manager (planning) + Executor (action) workflows
 """
 
 import logging
 from typing import List
 
+import llama_index.core
 from llama_index.core.llms.llm import LLM
-from llama_index.core.workflow import step, StartEvent, StopEvent, Workflow, Context
+from llama_index.core.workflow import Context, StartEvent, StopEvent, Workflow, step
 from llama_index.core.workflow.handler import WorkflowHandler
-from droidrun.agent.droid.events import *
+from workflows.events import Event
+
+from droidrun.config_manager.config_manager import VisionConfig
+
 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, MacroEvent, RecordUIStateEvent
-from droidrun.agent.common.default import MockWorkflow
+from droidrun.agent.common.events import MacroEvent, RecordUIStateEvent, ScreenshotEvent
 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
+from droidrun.agent.context.task_manager import Task, TaskManager
+from droidrun.agent.droid.events import (
+    CodeActExecuteEvent,
+    CodeActResultEvent,
+    DroidAgentState,
+    ExecutorInputEvent,
+    ExecutorResultEvent,
+    FinalizeEvent,
+    ManagerInputEvent,
+    ManagerPlanEvent,
+)
+from droidrun.agent.executor import ExecutorAgent
+from droidrun.agent.manager import ManagerAgent
+from droidrun.agent.utils.tools import ATOMIC_ACTION_SIGNATURES
+from droidrun.agent.utils.trajectory import Trajectory
 from droidrun.telemetry import (
+    DroidAgentFinalizeEvent,
+    DroidAgentInitEvent,
     capture,
     flush,
-    DroidAgentInitEvent,
-    DroidAgentFinalizeEvent,
 )
+from droidrun.telemetry.phoenix import arize_phoenix_callback_handler
+from droidrun.tools import Tools
 
 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.
+    A wrapper class that coordinates between agents to achieve a user's goal.
+
+    Reasoning modes:
+    - reasoning=False: Uses CodeActAgent directly for immediate execution
+    - reasoning=True: Uses ManagerAgent (planning) + ExecutorAgent (actions)
     """
 
     @staticmethod
@@ -63,18 +85,18 @@ class DroidAgent(Workflow):
     def __init__(
         self,
         goal: str,
-        llm: LLM,
+        llms: dict[str, LLM] | LLM,
         tools: Tools,
-        personas: List[AgentPersona] = [DEFAULT],
+        personas: List[AgentPersona] = [DEFAULT],  # noqa: B006
         max_steps: int = 15,
         timeout: int = 1000,
-        vision: bool = False,
+        vision: "VisionConfig | dict | bool" = False,
         reasoning: bool = False,
-        reflection: bool = False,
         enable_tracing: bool = False,
         debug: bool = False,
         save_trajectories: str = "none",
         excluded_tools: List[str] = None,
+        custom_tools: dict = None,
         *args,
         **kwargs,
     ):
@@ -86,42 +108,97 @@ class DroidAgent(Workflow):
             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)
+            reasoning: Whether to use Manager+Executor 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
             save_trajectories: Trajectory saving level. Can be:
                 - "none" (no saving)
                 - "step" (save per step)
                 - "action" (save per action)
+            custom_tools: Dictionary of custom tools in ATOMIC_ACTION_SIGNATURES format:
+                {
+                    "tool_name": {
+                        "arguments": ["arg1", "arg2"],
+                        "description": "Tool description with usage example",
+                        "function": callable
+                    }
+                }
             **kwargs: Additional keyword arguments to pass to the agents
         """
         self.user_id = kwargs.pop("user_id", None)
-        super().__init__(timeout=timeout, *args, **kwargs)
+        super().__init__(timeout=timeout, *args, **kwargs)  # noqa: B026
         # 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")
+                handler = arize_phoenix_callback_handler()
+                llama_index.core.global_handler = handler
                 logger.info("🔍 Arize Phoenix tracing enabled globally")
             except ImportError:
-                logger.warning("⚠️ Arize Phoenix package not found, tracing disabled")
+                logger.warning(
+                    "⚠️  Arize Phoenix is not installed.\n"
+                    "    To enable Phoenix integration, install with:\n"
+                    "    • If installed via tool: `uv tool install droidrun[phoenix]`"
+                    "    • If installed via pip: `uv pip install droidrun[phoenix]`\n"
+                )
                 enable_tracing = False
 
         self.goal = goal
-        self.llm = llm
-        self.vision = vision
         self.max_steps = max_steps
         self.max_codeact_steps = max_steps
         self.timeout = timeout
         self.reasoning = reasoning
-        self.reflection = reflection
         self.debug = debug
+        self.custom_tools = custom_tools or {}
+
+        # ====================================================================
+        # Handle LLM parameter - support both dict and single LLM
+        # ====================================================================
+        if isinstance(llms, dict):
+            self.manager_llm = llms.get('manager')
+            self.executor_llm = llms.get('executor')
+            self.codeact_llm = llms.get('codeact')
+            self.text_manipulator_llm = llms.get('text_manipulator')
+            self.app_opener_llm = llms.get('app_opener')
+            
+            # Validate required LLMs are present
+            if reasoning and (not self.manager_llm or not self.executor_llm):
+                raise ValueError("When reasoning=True, 'manager' and 'executor' LLMs must be provided in llms dict")
+            if not self.codeact_llm:
+                raise ValueError("'codeact' LLM must be provided in llms dict")
+            
+            logger.info("📚 Using agent-specific LLMs from dictionary")
+        else:
+            # single LLM for all agents
+            logger.info("📚 Using single LLM for all agents (backward compatibility mode)")
+            self.manager_llm = llms
+            self.executor_llm = llms
+            self.codeact_llm = llms
+            self.text_manipulator_llm = llms
+            self.app_opener_llm = llms
+
+        # ====================================================================
+        # Handle vision parameter - support VisionConfig, dict, or bool
+        # ====================================================================
+        if isinstance(vision, VisionConfig):
+            self.vision_config = vision
+        elif isinstance(vision, dict):
+            self.vision_config = VisionConfig.from_dict(vision)
+        elif isinstance(vision, bool):
+            # Backward compatibility: single bool for all agents
+            logger.info(f"👁️  Using vision={vision} for all agents (backward compatibility mode)")
+            self.vision_config = VisionConfig(manager=vision, executor=vision, codeact=vision)
+        else:
+            raise TypeError(f"vision must be VisionConfig, dict, or bool, got {type(vision)}")
+
+        # Store individual vision flags for easy access
+        self.manager_vision = self.vision_config.manager
+        self.executor_vision = self.vision_config.executor
+        self.codeact_vision = self.vision_config.codeact
+
 
         self.event_counter = 0
         # Handle backward compatibility: bool -> str mapping
@@ -148,43 +225,64 @@ class DroidAgent(Workflow):
         logger.info("🤖 Initializing DroidAgent...")
         logger.info(f"💾 Trajectory saving level: {self.save_trajectories}")
 
-        self.tool_list = describe_tools(tools, excluded_tools)
         self.tools_instance = tools
 
         self.tools_instance.save_trajectories = self.save_trajectories
 
+        # Create shared state instance for Manager/Executor workflows
+        self.shared_state = DroidAgentState(
+            instruction=goal,
+            err_to_manager_thresh=2
+        )
+
         if self.reasoning:
-            logger.info("📝 Initializing Planner Agent...")
-            self.planner_agent = PlannerAgent(
-                goal=goal,
-                llm=llm,
-                vision=vision,
+            logger.info("📝 Initializing Manager and Executor Agents...")
+            self.manager_agent = ManagerAgent(
+                llm=self.manager_llm,
+                vision=self.manager_vision,
                 personas=personas,
-                task_manager=self.task_manager,
                 tools_instance=tools,
+                shared_state=self.shared_state,
+                custom_tools=self.custom_tools,
+                timeout=timeout,
+                debug=debug,
+            )
+            self.executor_agent = ExecutorAgent(
+                llm=self.executor_llm,
+                vision=self.executor_vision,
+                tools_instance=tools,
+                shared_state=self.shared_state,
+                persona=None,  # Need to figure this out
+                custom_tools=self.custom_tools,
                 timeout=timeout,
                 debug=debug,
             )
             self.max_codeact_steps = 5
 
-            if self.reflection:
-                self.reflector = Reflector(llm=llm, debug=debug)
+
+            # Keep planner_agent for backward compatibility (can be removed later)
+            self.planner_agent = None
 
         else:
-            logger.debug("🚫 Planning disabled - will execute tasks directly with CodeActAgent")
+            logger.debug("🚫 Reasoning disabled - will execute tasks directly with CodeActAgent")
+            self.manager_agent = None
+            self.executor_agent = None
             self.planner_agent = None
 
+        # Get tool names from ATOMIC_ACTION_SIGNATURES for telemetry
+        atomic_tools = list(ATOMIC_ACTION_SIGNATURES.keys())
+
         capture(
+            # TODO: do proper telemetry instead of this ductaped crap
             DroidAgentInitEvent(
                 goal=goal,
-                llm=llm.class_name(),
-                tools=",".join(self.tool_list),
+                llm=self.llm.class_name(),
+                tools=",".join(atomic_tools + ["remember", "complete"]),
                 personas=",".join([p.name for p in personas]),
                 max_steps=max_steps,
                 timeout=timeout,
-                vision=vision,
+                vision=self.vision,
                 reasoning=reasoning,
-                reflection=reflection,
                 enable_tracing=enable_tracing,
                 debug=debug,
                 save_trajectories=save_trajectories,
@@ -200,6 +298,35 @@ class DroidAgent(Workflow):
         """
         return super().run(*args, **kwargs)
 
+    def _create_finalize_event(
+        self,
+        success: bool,
+        reason: str,
+        output: str
+    ) -> FinalizeEvent:
+        """
+        Single source of truth for creating FinalizeEvent.
+
+        This helper ensures all FinalizeEvent creation is consistent
+        across the workflow.
+
+        Args:
+            success: Whether the task succeeded
+            reason: Reason for completion (deprecated, use output)
+            output: Output message
+
+        Returns:
+            FinalizeEvent ready to be returned
+        """
+        return FinalizeEvent(
+            success=success,
+            reason=reason,
+            output=output,
+            task=[], # TODO: use the final plan as the tasks and the goal as task
+            tasks=[],
+            steps=self.step_counter
+        )
+
     @step
     async def execute_task(self, ctx: Context, ev: CodeActExecuteEvent) -> CodeActResultEvent:
         """
@@ -212,19 +339,18 @@ class DroidAgent(Workflow):
             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,
+                llm=self.codeact_llm,
                 persona=persona,
-                vision=self.vision,
+                vision=self.codeact_vision,
                 max_steps=self.max_codeact_steps,
-                all_tools_list=self.tool_list,
                 tools_instance=self.tools_instance,
+                custom_tools=self.custom_tools,
                 debug=self.debug,
                 timeout=self.timeout,
             )
@@ -232,7 +358,6 @@ class DroidAgent(Workflow):
             handler = codeact_agent.run(
                 input=task.description,
                 remembered_info=self.tools_instance.memory,
-                reflection=reflection,
             )
 
             async for nested_ev in handler.stream_events():
@@ -266,32 +391,17 @@ class DroidAgent(Workflow):
     @step
     async def handle_codeact_execute(
         self, ctx: Context, ev: CodeActResultEvent
-    ) -> FinalizeEvent | ReflectionEvent | ReasoningLogicEvent:
+    ) -> FinalizeEvent:
         try:
             task = ev.task
-            if not self.reasoning:
-                return FinalizeEvent(
-                    success=ev.success,
-                    reason=ev.reason,
-                    output=ev.reason,
-                    task=[task],
-                    tasks=[task],
-                    steps=ev.steps,
-                )
-
-            if self.reflection and ev.success:
-                return ReflectionEvent(task=task)
-
-            # Reasoning is enabled but reflection is disabled.
-            # Success: mark complete and proceed to next step in reasoning loop.
-            # Failure: mark failed and trigger planner immediately without advancing to the next queued task.
-            if ev.success:
-                self.task_manager.complete_task(task, message=ev.reason)
-                return ReasoningLogicEvent()
-            else:
-                self.task_manager.fail_task(task, failure_reason=ev.reason)
-                return ReasoningLogicEvent(force_planning=True)
-
+            return FinalizeEvent(
+                success=ev.success,
+                reason=ev.reason,
+                output=ev.reason,
+                task=[task],
+                tasks=[task],
+                steps=ev.steps,
+            )
         except Exception as e:
             logger.error(f"❌ Error during DroidAgent execution: {e}")
             if self.debug:
@@ -309,141 +419,178 @@ class DroidAgent(Workflow):
             )
 
     @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
-        )
+    async def start_handler(
+        self, ctx: Context, ev: StartEvent
+    ) -> CodeActExecuteEvent | ManagerInputEvent:
+        """
+        Main execution loop that coordinates between planning and execution.
 
-        if reflection.goal_achieved:
-            self.task_manager.complete_task(task)
-            return ReasoningLogicEvent()
+        Returns:
+            Event to trigger next step based on reasoning mode
+        """
+        logger.info(f"🚀 Running DroidAgent to achieve goal: {self.goal}")
+        ctx.write_event_to_stream(ev)
 
-        else:
-            self.task_manager.fail_task(task)
-            return ReasoningLogicEvent(reflection=reflection)
+        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)
+
+        # Reasoning mode - state already initialized in __init__, start with Manager
+        logger.info("🧠 Reasoning mode - initializing Manager/Executor workflow")
+        return ManagerInputEvent()
+
+    # ========================================================================
+    # Manager/Executor Workflow Steps
+    # ========================================================================
 
     @step
-    async def handle_reasoning_logic(
+    async def run_manager(
         self,
         ctx: Context,
-        ev: ReasoningLogicEvent,
-    ) -> FinalizeEvent | CodeActExecuteEvent:
-        try:
-            if self.step_counter >= self.max_steps:
-                output = f"Reached maximum number of steps ({self.max_steps})"
-                tasks = self.task_manager.get_task_history()
-                return FinalizeEvent(
-                    success=False,
-                    reason=output,
-                    output=output,
-                    task=tasks,
-                    tasks=tasks,
-                    steps=self.step_counter,
-                )
-            self.step_counter += 1
+        ev: ManagerInputEvent
+    ) -> ManagerPlanEvent | FinalizeEvent:
+        """
+        Run Manager planning phase.
 
-            if ev.reflection:
-                handler = self.planner_agent.run(
-                    remembered_info=self.tools_instance.memory, reflection=ev.reflection
-                )
-            else:
-                if not ev.force_planning and self.task_iter:
-                    try:
-                        task = next(self.task_iter)
-                        return CodeActExecuteEvent(task=task, reflection=None)
-                    except StopIteration as e:
-                        logger.info("Planning next steps...")
+        Pre-flight checks for termination before running manager.
+        The Manager analyzes current state and creates a plan with subgoals.
+        """
+        # Check if we've reached the maximum number of steps
+        if self.step_counter >= self.max_steps:
+            logger.warning(f"⚠️ Reached maximum steps ({self.max_steps})")
+            return self._create_finalize_event(
+                success=False,
+                reason=f"Reached maximum steps ({self.max_steps})",
+                output=f"Reached maximum steps ({self.max_steps})"
+            )
 
-                logger.debug(f"Planning step {self.step_counter}/{self.max_steps}")
+        # Continue with Manager execution
+        logger.info(f"📋 Running Manager for planning... (step {self.step_counter}/{self.max_steps})")
 
-                handler = self.planner_agent.run(
-                    remembered_info=self.tools_instance.memory, reflection=None
-                )
+        # Run Manager workflow
+        handler = self.manager_agent.run()
 
-            async for nested_ev in handler.stream_events():
-                self.handle_stream_event(nested_ev, ctx)
+        # Stream nested events
+        async for nested_ev in handler.stream_events():
+            self.handle_stream_event(nested_ev, ctx)
 
-            result = await handler
+        result = await handler
 
-            self.tasks = self.task_manager.get_all_tasks()
-            self.task_iter = iter(self.tasks)
+        # Manager already updated shared_state, just return event with results
+        return ManagerPlanEvent(
+            plan=result["plan"],
+            current_subgoal=result["current_subgoal"],
+            completed_plan=result["completed_plan"],
+            thought=result["thought"],
+            manager_answer=result.get("manager_answer", "")
+        )
 
-            if self.task_manager.goal_completed:
-                logger.info(f"✅ Goal completed: {self.task_manager.message}")
-                tasks = self.task_manager.get_task_history()
-                return FinalizeEvent(
-                    success=True,
-                    reason=self.task_manager.message,
-                    output=self.task_manager.message,
-                    task=tasks,
-                    tasks=tasks,
-                    steps=self.step_counter,
-                )
-            if not self.tasks:
-                logger.warning("No tasks generated by planner")
-                output = "Planner did not generate any tasks"
-                tasks = self.task_manager.get_task_history()
-                return FinalizeEvent(
-                    success=False,
-                    reason=output,
-                    output=output,
-                    task=tasks,
-                    tasks=tasks,
-                    steps=self.step_counter,
-                )
+    @step
+    async def handle_manager_plan(
+        self,
+        ctx: Context,
+        ev: ManagerPlanEvent
+    ) -> ExecutorInputEvent | FinalizeEvent:
+        """
+        Process Manager output and decide next step.
 
-            return CodeActExecuteEvent(task=next(self.task_iter), reflection=None)
+        Checks if task is complete or if Executor should take action.
+        """
+        # Check for answer-type termination
+        if ev.manager_answer.strip():
+            logger.info(f"💬 Manager provided answer: {ev.manager_answer}")
+            self.shared_state.progress_status = f"Answer: {ev.manager_answer}"
+
+            return self._create_finalize_event(
+                success=True,
+                reason=ev.manager_answer,
+                output=ev.manager_answer
+            )
 
-        except Exception as e:
-            logger.error(f"❌ Error during DroidAgent execution: {e}")
-            if self.debug:
-                import traceback
+        # Continue to Executor with current subgoal
+        logger.info(f"▶️  Proceeding to Executor with subgoal: {ev.current_subgoal}")
+        return ExecutorInputEvent(current_subgoal=ev.current_subgoal)
 
-                logger.error(traceback.format_exc())
-            tasks = self.task_manager.get_task_history()
-            return FinalizeEvent(
-                success=False,
-                reason=str(e),
-                output=str(e),
-                task=tasks,
-                tasks=tasks,
-                steps=self.step_counter,
-            )
+    @step
+    async def run_executor(
+        self,
+        ctx: Context,
+        ev: ExecutorInputEvent
+    ) -> ExecutorResultEvent:
+        """
+        Run Executor action phase.
+
+        The Executor selects and executes a specific action for the current subgoal.
+        """
+        logger.info("⚡ Running Executor for action...")
+
+        # Run Executor workflow (Executor will update shared_state directly)
+        handler = self.executor_agent.run(subgoal=ev.current_subgoal)
+
+        # Stream nested events
+        async for nested_ev in handler.stream_events():
+            self.handle_stream_event(nested_ev, ctx)
+
+        result = await handler
+
+        # Update coordination state after execution
+        self.shared_state.action_history.append(result["action"])
+        self.shared_state.summary_history.append(result["summary"])
+        self.shared_state.action_outcomes.append(result["outcome"])
+        self.shared_state.error_descriptions.append(result["error"])
+        self.shared_state.last_action = result["action"]
+        self.shared_state.last_summary = result["summary"]
+        self.shared_state.last_action_thought = result.get("thought", "")
+        self.shared_state.action_pool.append(result["action_json"])
+        self.shared_state.progress_status = self.shared_state.completed_plan
+
+        return ExecutorResultEvent(
+            action=result["action"],
+            outcome=result["outcome"],
+            error=result["error"],
+            summary=result["summary"]
+        )
 
     @step
-    async def start_handler(
-        self, ctx: Context, ev: StartEvent
-    ) -> CodeActExecuteEvent | ReasoningLogicEvent:
+    async def handle_executor_result(
+        self,
+        ctx: Context,
+        ev: ExecutorResultEvent
+    ) -> ManagerInputEvent:
         """
-        Main execution loop that coordinates between planning and execution.
+        Process Executor result and continue.
 
-        Returns:
-            Dict containing the execution result
+        Checks for error escalation and loops back to Manager.
+        Note: Max steps check is now done in run_manager pre-flight.
         """
-        logger.info(f"🚀 Running DroidAgent to achieve goal: {self.goal}")
-        ctx.write_event_to_stream(ev)
+        # Check error escalation
+        err_thresh = self.shared_state.err_to_manager_thresh
 
-        self.step_counter = 0
-        self.retry_counter = 0
+        if len(self.shared_state.action_outcomes) >= err_thresh:
+            latest = self.shared_state.action_outcomes[-err_thresh:]
+            error_count = sum(1 for o in latest if not o)
+            if error_count == err_thresh:
+                logger.warning(f"⚠️ Error escalation: {err_thresh} consecutive errors")
+                self.shared_state.error_flag_plan = True
 
-        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",
-            )
+        self.step_counter += 1
+        logger.info(f"🔄 Step {self.step_counter}/{self.max_steps} complete, looping to Manager")
 
-            return CodeActExecuteEvent(task=task, reflection=None)
+        # Always loop back to Manager (it will check max steps in pre-flight)
+        return ManagerInputEvent()
 
-        return ReasoningLogicEvent()
+    # ========================================================================
+    # End Manager/Executor Workflow Steps
+    # ========================================================================
 
     @step
     async def finalize(self, ctx: Context, ev: FinalizeEvent) -> StopEvent: