droidrun 0.3.5__py3-none-any.whl → 0.3.6__py3-none-any.whl

droidrun/agent/codeact/codeact_agent.py

@@ -18,7 +18,8 @@ from droidrun.agent.codeact.events import (
     TaskThinkingEvent,
     EpisodicMemoryEvent,
 )
-from droidrun.agent.common.events import ScreenshotEvent
+from droidrun.agent.common.events import ScreenshotEvent, RecordUIStateEvent
+from droidrun.agent.usage import get_usage_from_response
 from droidrun.agent.utils import chat_utils
 from droidrun.agent.utils.executer import SimpleCodeExecutor
 from droidrun.agent.codeact.prompts import (
@@ -182,6 +183,7 @@ class CodeActAgent(Workflow):
                 try:
                     state = self.tools.get_state()
                     await ctx.set("ui_state", state["a11y_tree"])
+                    ctx.write_event_to_stream(RecordUIStateEvent(ui_state=state["a11y_tree"]))
                     chat_history = await chat_utils.add_ui_text_block(
                         state["a11y_tree"], chat_history
                     )
@@ -202,11 +204,17 @@ class CodeActAgent(Workflow):
                 success=False, reason="LLM response is None. This is a critical error."
             )
 
+        try:
+            usage = get_usage_from_response(self.llm.class_name(), response)
+        except Exception as e:
+            logger.warning(f"Could not get llm usage from response: {e}")
+            usage = None
+
         await self.chat_memory.aput(response.message)
 
         code, thoughts = chat_utils.extract_code_and_thought(response.message.content)
 
-        event = TaskThinkingEvent(thoughts=thoughts, code=code)
+        event = TaskThinkingEvent(thoughts=thoughts, code=code, usage=usage)
         ctx.write_event_to_stream(event)
         return event
 
@@ -255,6 +263,10 @@ class CodeActAgent(Workflow):
             for screenshot in screenshots[:-1]: # the last screenshot will be captured by next step
                 ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
 
+            ui_states = result['ui_states']
+            for ui_state in ui_states[:-1]:
+                ctx.write_event_to_stream(RecordUIStateEvent(ui_state=ui_state['a11y_tree']))
+
             if self.tools.finished == True:
                 logger.debug("  - Task completed.")
                 event = TaskEndEvent(
@@ -311,7 +323,8 @@ class CodeActAgent(Workflow):
         await ctx.set("chat_memory", self.chat_memory)
         
         # Add final state observation to episodic memory
-        await self._add_final_state_observation(ctx)
+        if self.vision:
+            await self._add_final_state_observation(ctx)
         
         result = {}
         result.update(

droidrun/agent/codeact/events.py

@@ -1,6 +1,8 @@
 from llama_index.core.llms import ChatMessage
 from llama_index.core.workflow import Event
 from typing import Optional
+
+from droidrun.agent.usage import UsageResult
 from ..context.episodic_memory import EpisodicMemory
 
 class TaskInputEvent(Event):
@@ -11,6 +13,7 @@ class TaskInputEvent(Event):
 class TaskThinkingEvent(Event):
     thoughts: Optional[str] = None
     code: Optional[str] = None  
+    usage: Optional[UsageResult] = None
 
 class TaskExecutionEvent(Event):
     code: str

droidrun/agent/common/events.py

@@ -1,4 +1,5 @@
 from llama_index.core.workflow import Event
+from typing import Dict, Any
 
 class ScreenshotEvent(Event):
     screenshot: bytes
@@ -44,4 +45,7 @@ class KeyPressActionEvent(MacroEvent):
 class StartAppEvent(MacroEvent):
     """"Event for starting an app"""
     package: str
-    activity: str = None
+    activity: str = None
+
+class RecordUIStateEvent(Event):
+    ui_state: list[Dict[str, Any]]

droidrun/agent/droid/droid_agent.py

@@ -16,23 +16,28 @@ 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
+from droidrun.agent.common.events import ScreenshotEvent, MacroEvent, RecordUIStateEvent
 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
-from droidrun.telemetry import capture, flush, DroidAgentInitEvent, DroidAgentFinalizeEvent
-
+from droidrun.telemetry import (
+    capture,
+    flush,
+    DroidAgentInitEvent,
+    DroidAgentFinalizeEvent,
+)
 
 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 PlannerAgent (creates plans) and
+        CodeActAgent (executes tasks) to achieve a user's goal.
     """
-    
+
     @staticmethod
     def _configure_default_logging(debug: bool = False):
         """
@@ -43,20 +48,20 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
         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, 
+        self,
         goal: str,
         llm: LLM,
         tools: Tools,
@@ -71,17 +76,17 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
         save_trajectories: str = "none",
         excluded_tools: List[str] = None,
         *args,
-        **kwargs
+        **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) 
+            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
@@ -93,14 +98,15 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             **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)
         # 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:
@@ -125,27 +131,27 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             # Validate string values
             valid_values = ["none", "step", "action"]
             if save_trajectories not in valid_values:
-                logger.warning(f"Invalid save_trajectories value: {save_trajectories}. Using 'none' instead.")
+                logger.warning(
+                    f"Invalid save_trajectories value: {save_trajectories}. Using 'none' instead."
+                )
                 self.save_trajectories = "none"
             else:
                 self.save_trajectories = save_trajectories
-        
+
         self.trajectory = Trajectory(goal=goal)
         self.task_manager = TaskManager()
         self.task_iter = None
 
-        
         self.cim = ContextInjectionManager(personas=personas)
         self.current_episodic_memory = None
 
         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
 
+        self.tools_instance.save_trajectories = self.save_trajectories
 
         if self.reasoning:
             logger.info("📝 Initializing Planner Agent...")
@@ -157,14 +163,14 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
                 task_manager=self.task_manager,
                 tools_instance=tools,
                 timeout=timeout,
-                debug=debug
+                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
@@ -184,10 +190,9 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
                 debug=debug,
                 save_trajectories=save_trajectories,
             ),
-            self.user_id
+            self.user_id,
         )
 
-        
         logger.info("✅ DroidAgent initialized successfully.")
 
     def run(self, *args, **kwargs) -> WorkflowHandler:
@@ -195,19 +200,15 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
         Run the DroidAgent workflow.
         """
         return super().run(*args, **kwargs)
-    
+
     @step
-    async def execute_task(
-        self,
-        ctx: Context,
-        ev: CodeActExecuteEvent
-        ) -> CodeActResultEvent:
+    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)
         """
@@ -232,34 +233,53 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             handler = codeact_agent.run(
                 input=task.description,
                 remembered_info=self.tools_instance.memory,
-                reflection=reflection
+                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"])
+                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"])
-                
+                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=[])
-    
+
     @step
-    async def handle_codeact_execute(self, ctx: Context, ev: CodeActResultEvent) -> FinalizeEvent | ReflectionEvent | ReasoningLogicEvent:
+    async def handle_codeact_execute(
+        self, ctx: Context, ev: CodeActResultEvent
+    ) -> FinalizeEvent | ReflectionEvent | ReasoningLogicEvent:
         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)
-            
+                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)
 
@@ -277,51 +297,64 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             logger.error(f"❌ Error during DroidAgent execution: {e}")
             if self.debug:
                 import traceback
+
                 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)
-        
+            return FinalizeEvent(
+                success=False,
+                reason=str(e),
+                output=str(e),
+                task=tasks,
+                tasks=tasks,
+                steps=self.step_counter,
+            )
 
     @step
     async def reflect(
-        self,
-        ctx: Context,
-        ev: ReflectionEvent
-        ) -> ReasoningLogicEvent | CodeActExecuteEvent:
-
-
+        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)
+
+        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:
+        planner_agent: Workflow = MockWorkflow(),
+    ) -> 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)
+                return FinalizeEvent(
+                    success=False,
+                    reason=output,
+                    output=output,
+                    task=tasks,
+                    tasks=tasks,
+                    steps=self.step_counter,
+                )
             self.step_counter += 1
 
             if ev.reflection:
-                handler = planner_agent.run(remembered_info=self.tools_instance.memory, reflection=ev.reflection)
+                handler = planner_agent.run(
+                    remembered_info=self.tools_instance.memory, reflection=ev.reflection
+                )
             else:
                 if not ev.force_planning and self.task_iter:
                     try:
@@ -332,7 +365,9 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
 
                 logger.debug(f"Planning step {self.step_counter}/{self.max_steps}")
 
-                handler = planner_agent.run(remembered_info=self.tools_instance.memory, reflection=None)
+                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)
@@ -345,51 +380,73 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             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)
+                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)
-            
+                return FinalizeEvent(
+                    success=False,
+                    reason=output,
+                    output=output,
+                    task=tasks,
+                    tasks=tasks,
+                    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())
             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)
-    
+            return FinalizeEvent(
+                success=False,
+                reason=str(e),
+                output=str(e),
+                task=tasks,
+                tasks=tasks,
+                steps=self.step_counter,
+            )
 
     @step
-    async def start_handler(self, ctx: Context, ev: StartEvent) -> CodeActExecuteEvent | ReasoningLogicEvent:
+    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}")
         ctx.write_event_to_stream(ev)
-        
+
         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"
+                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)
@@ -400,7 +457,7 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
                 output=ev.output,
                 steps=ev.steps,
             ),
-            self.user_id
+            self.user_id,
         )
         flush()
 
@@ -416,24 +473,20 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             self.trajectory.save_trajectory()
 
         return StopEvent(result)
-    
-    def handle_stream_event(self, ev: Event, ctx: Context):
 
+    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)
             elif isinstance(ev, MacroEvent):
                 self.trajectory.macro.append(ev)
+            elif isinstance(ev, RecordUIStateEvent):
+                self.trajectory.ui_states.append(ev.ui_state)
             else:
                 self.trajectory.events.append(ev)
-
-
-

droidrun/agent/planner/events.py

@@ -2,6 +2,7 @@ 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
+from droidrun.agent.usage import UsageResult
 
 class PlanInputEvent(Event):
     input: list[ChatMessage]
@@ -10,6 +11,7 @@ class PlanInputEvent(Event):
 class PlanThinkingEvent(Event):
     thoughts: Optional[str] = None
     code: Optional[str] = None  
+    usage: Optional[UsageResult] = None
 
 
 class PlanCreatedEvent(Event):

droidrun/agent/planner/planner_agent.py

@@ -13,11 +13,12 @@ from llama_index.core.llms.llm import LLM
 from llama_index.core.workflow import Workflow, StartEvent, StopEvent, Context, step
 from llama_index.core.memory import Memory
 from llama_index.core.llms.llm import LLM
+from droidrun.agent.usage import get_usage_from_response
 from droidrun.agent.utils.executer import SimpleCodeExecutor
 from droidrun.agent.utils import chat_utils
 from droidrun.agent.context.task_manager import TaskManager
 from droidrun.tools import Tools
-from droidrun.agent.common.events import ScreenshotEvent
+from droidrun.agent.common.events import ScreenshotEvent, RecordUIStateEvent
 from droidrun.agent.planner.events import (
     PlanInputEvent,
     PlanCreatedEvent,
@@ -130,16 +131,16 @@ class PlannerAgent(Workflow):
         self.steps_counter += 1
         logger.info(f"🧠 Thinking about how to plan the goal...")
 
-        # if vision is disabled, screenshot should save to trajectory
-        screenshot = (self.tools_instance.take_screenshot())[1]
-        ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
         if self.vision:
+            screenshot = (self.tools_instance.take_screenshot())[1]
+            ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
             await ctx.set("screenshot", screenshot)
 
         try:
             state = self.tools_instance.get_state()
             await ctx.set("ui_state", state["a11y_tree"])
             await ctx.set("phone_state", state["phone_state"])
+            ctx.write_event_to_stream(RecordUIStateEvent(ui_state=state["a11y_tree"]))
         except Exception as e:
             logger.warning(f"⚠️ Error retrieving state from the connected device. Is the Accessibility Service enabled?")
 
@@ -148,11 +149,16 @@ class PlannerAgent(Workflow):
         await ctx.set("reflection", self.reflection)
 
         response = await self._get_llm_response(ctx, chat_history)
+        try:
+            usage = get_usage_from_response(self.llm.class_name(), response)
+        except Exception as e:
+            logger.warning(f"Could not get llm usage from response: {e}")
+            usage = None
         await self.chat_memory.aput(response.message)
 
         code, thoughts = chat_utils.extract_code_and_thought(response.message.content)
 
-        event = PlanThinkingEvent(thoughts=thoughts, code=code)
+        event = PlanThinkingEvent(thoughts=thoughts, code=code, usage=usage)
         ctx.write_event_to_stream(event)
         return event
 
@@ -174,6 +180,10 @@ class PlannerAgent(Workflow):
                 screenshots = result['screenshots']
                 for screenshot in screenshots[:-1]: # the last screenshot will be captured by next step
                     ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
+                
+                ui_states = result['ui_states']
+                for ui_state in ui_states[:-1]:
+                    ctx.write_event_to_stream(RecordUIStateEvent(ui_state=ui_state['a11y_tree']))
 
                 await self.chat_memory.aput(
                     ChatMessage(