droidrun 0.3.8__py3-none-any.whl → 0.3.10.dev2__py3-none-any.whl

droidrun/__init__.py

@@ -5,13 +5,12 @@ DroidRun - A framework for controlling Android devices through LLM agents.
 __version__ = "0.3.0"
 
 # Import main classes for easier access
-from droidrun.agent.utils.llm_picker import load_llm
-from droidrun.tools import Tools, AdbTools, IOSTools
 from droidrun.agent.droid import DroidAgent
+from droidrun.agent.utils.llm_picker import load_llm
 
 # Import macro functionality
 from droidrun.macro import MacroPlayer, replay_macro_file, replay_macro_folder
-
+from droidrun.tools import AdbTools, IOSTools, Tools
 
 # Make main components available at package level
 __all__ = [

droidrun/__main__.py

@@ -4,4 +4,4 @@ DroidRun main entry point
 from droidrun.cli.main import cli
 
 if __name__ == '__main__':
-    cli()
+    cli()

droidrun/agent/__init__.py

@@ -3,4 +3,4 @@
 #logger = logging.getLogger("droidrun")
 #logger.propagate = False  # Don't send to root logger
 #logger.handlers = []      # No handlers by default
-#logger.setLevel(logging.INFO)  # Or WARNING
+#logger.setLevel(logging.INFO)  # Or WARNING

droidrun/agent/codeact/__init__.py

@@ -1,8 +1,5 @@
 from droidrun.agent.codeact.codeact_agent import CodeActAgent
-from droidrun.agent.codeact.prompts import (
-    DEFAULT_CODE_ACT_USER_PROMPT,
-    DEFAULT_NO_THOUGHTS_PROMPT
-)
+from droidrun.agent.codeact.prompts import DEFAULT_CODE_ACT_USER_PROMPT, DEFAULT_NO_THOUGHTS_PROMPT
 
 __all__ = [
     "CodeActAgent",

droidrun/agent/codeact/codeact_agent.py

@@ -1,36 +1,37 @@
+import asyncio
+import json
 import logging
 import re
 import time
-import asyncio
-import json
-import os
-from typing import List, Optional, Tuple, Union
+from typing import List, Union
+
 from llama_index.core.base.llms.types import ChatMessage, ChatResponse
-from llama_index.core.prompts import PromptTemplate
 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.prompts import PromptTemplate
+from llama_index.core.workflow import Context, StartEvent, StopEvent, Workflow, step
+
 from droidrun.agent.codeact.events import (
-    TaskInputEvent,
+    EpisodicMemoryEvent,
     TaskEndEvent,
     TaskExecutionEvent,
     TaskExecutionResultEvent,
+    TaskInputEvent,
     TaskThinkingEvent,
-    EpisodicMemoryEvent,
 )
-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 (
     DEFAULT_CODE_ACT_USER_PROMPT,
     DEFAULT_NO_THOUGHTS_PROMPT,
 )
-
+from droidrun.agent.common.constants import LLM_HISTORY_LIMIT
+from droidrun.agent.common.events import RecordUIStateEvent, ScreenshotEvent
+from droidrun.agent.context.agent_persona import AgentPersona
 from droidrun.agent.context.episodic_memory import EpisodicMemory, EpisodicMemoryStep
+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.utils.tools import ATOMIC_ACTION_SIGNATURES, get_atomic_tool_descriptions, build_custom_tool_descriptions
 from droidrun.tools import Tools
-from typing import Optional, Dict, Tuple, List, Any, Callable
-from droidrun.agent.context.agent_persona import AgentPersona
 
 logger = logging.getLogger("droidrun")
 
@@ -48,8 +49,8 @@ class CodeActAgent(Workflow):
         persona: AgentPersona,
         vision: bool,
         tools_instance: "Tools",
-        all_tools_list: Dict[str, Callable[..., Any]],
         max_steps: int = 5,
+        custom_tools: dict = None,
         debug: bool = False,
         *args,
         **kwargs,
@@ -77,13 +78,40 @@ class CodeActAgent(Workflow):
 
         self.tools = tools_instance
 
+        # Merge custom_tools with ATOMIC_ACTION_SIGNATURES
+        # Custom tools are treated the same as atomic actions by CodeAct
+        merged_signatures = {**ATOMIC_ACTION_SIGNATURES, **(custom_tools or {})}
+
+        # Build tool_list from merged signatures
         self.tool_list = {}
+        for action_name, signature in merged_signatures.items():
+            func = signature["function"]
+            # Create bound function (curry tools_instance as first argument)
+            # Handle both sync and async functions
+            if asyncio.iscoroutinefunction(func):
+                async def make_async_bound(f, ti):
+                    async def bound_func(*args, **kwargs):
+                        return await f(ti, *args, **kwargs)
+                    return bound_func
+                self.tool_list[action_name] = asyncio.run(make_async_bound(func, tools_instance))
+            else:
+                self.tool_list[action_name] = lambda *args, f=func, ti=tools_instance: f(ti, *args)
 
-        for tool_name in persona.allowed_tools:
-            if tool_name in all_tools_list:
-                self.tool_list[tool_name] = all_tools_list[tool_name]
+        # Add non-atomic tools (remember, complete) from tools_instance
+        self.tool_list["remember"] = tools_instance.remember
+        self.tool_list["complete"] = tools_instance.complete
 
-        self.tool_descriptions = chat_utils.parse_tool_descriptions(self.tool_list)
+        # Get tool descriptions from ATOMIC_ACTION_SIGNATURES and custom_tools
+        self.tool_descriptions = get_atomic_tool_descriptions()
+
+        # Add custom tool descriptions if provided
+        custom_descriptions = build_custom_tool_descriptions(custom_tools or {})
+        if custom_descriptions:
+            self.tool_descriptions += "\n" + custom_descriptions
+
+        # Add descriptions for remember/complete
+        self.tool_descriptions += "\n- remember(information: str): Remember information for later use"
+        self.tool_descriptions += "\n- complete(success: bool, reason: str): Mark task as complete"
 
         self.system_prompt_content = persona.system_prompt.format(
             tool_descriptions=self.tool_descriptions
@@ -109,7 +137,8 @@ class CodeActAgent(Workflow):
         """Prepare chat history from user input."""
         logger.info("💬 Preparing chat for task execution...")
 
-        self.chat_memory: Memory = await ctx.get(
+
+        self.chat_memory: Memory = await ctx.store.get(
             "chat_memory", default=Memory.from_defaults()
         )
 
@@ -135,7 +164,7 @@ class CodeActAgent(Workflow):
 
         await self.chat_memory.aput(self.user_message)
 
-        await ctx.set("chat_memory", self.chat_memory)
+        await ctx.store.set("chat_memory", self.chat_memory)
         input_messages = self.chat_memory.get_all()
         return TaskInputEvent(input=input_messages)
 
@@ -160,9 +189,9 @@ class CodeActAgent(Workflow):
         logger.info(f"🧠 Step {self.steps_counter}: Thinking...")
 
         model = self.llm.class_name()
-        
+
         if "remember" in self.tool_list and self.remembered_info:
-            await ctx.set("remembered_info", self.remembered_info)
+            await ctx.store.set("remembered_info", self.remembered_info)
             chat_history = await chat_utils.add_memory_block(self.remembered_info, chat_history)
 
         for context in self.required_context:
@@ -171,25 +200,25 @@ class CodeActAgent(Workflow):
                 screenshot = (self.tools.take_screenshot())[1]
                 ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
 
-                await ctx.set("screenshot", screenshot)
+                await ctx.store.set("screenshot", screenshot)
                 if model == "DeepSeek":
                     logger.warning(
                         "[yellow]DeepSeek doesnt support images. Disabling screenshots[/]"
                     )
-                elif self.vision == True: # if vision is enabled, add screenshot to chat history
+                elif self.vision: # if vision is enabled, add screenshot to chat history
                     chat_history = await chat_utils.add_screenshot_image_block(screenshot, chat_history)
 
             if context == "ui_state":
                 try:
                     state = self.tools.get_state()
-                    await ctx.set("ui_state", state["a11y_tree"])
+                    await ctx.store.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
                     )
                     chat_history = await chat_utils.add_phone_state_block(state["phone_state"], chat_history)
-                except Exception as e:
-                    logger.warning(f"⚠️ Error retrieving state from the connected device. Is the Accessibility Service enabled?")
+                except Exception:
+                    logger.warning("⚠️ Error retrieving state from the connected device. Is the Accessibility Service enabled?")
 
 
             if context == "packages":
@@ -252,7 +281,7 @@ class CodeActAgent(Workflow):
         """Execute the code and return the result."""
         code = ev.code
         assert code, "Code cannot be empty."
-        logger.info(f"⚡ Executing action...")
+        logger.info("⚡ Executing action...")
         logger.info(f"Code to execute:\n```python\n{code}\n```")
 
         try:
@@ -267,16 +296,16 @@ class CodeActAgent(Workflow):
             for ui_state in ui_states[:-1]:
                 ctx.write_event_to_stream(RecordUIStateEvent(ui_state=ui_state['a11y_tree']))
 
-            if self.tools.finished == True:
+            if self.tools.finished:
                 logger.debug("  - Task completed.")
                 event = TaskEndEvent(
                     success=self.tools.success, reason=self.tools.reason
                 )
                 ctx.write_event_to_stream(event)
                 return event
-            
+
             self.remembered_info = self.tools.memory
-            
+
             event = TaskExecutionResultEvent(output=str(result['output']))
             ctx.write_event_to_stream(event)
             return event
@@ -320,12 +349,12 @@ class CodeActAgent(Workflow):
     async def finalize(self, ev: TaskEndEvent, ctx: Context) -> StopEvent:
         """Finalize the workflow."""
         self.tools.finished = False
-        await ctx.set("chat_memory", self.chat_memory)
-        
+        await ctx.store.set("chat_memory", self.chat_memory)
+
         # Add final state observation to episodic memory
         if self.vision:
             await self._add_final_state_observation(ctx)
-        
+
         result = {}
         result.update(
             {
@@ -347,14 +376,15 @@ class CodeActAgent(Workflow):
         self, ctx: Context, chat_history: List[ChatMessage]
     ) -> ChatResponse | None:
         logger.debug("🔍 Getting LLM response...")
-        messages_to_send = [self.system_prompt] + chat_history
+        limited_history = self._limit_history(chat_history)
+        messages_to_send = [self.system_prompt] + limited_history
         messages_to_send = [chat_utils.message_copy(msg) for msg in messages_to_send]
         try:
             response = await self.llm.achat(messages=messages_to_send)
             logger.debug("🔍 Received LLM response.")
 
             filtered_chat_history = []
-            for msg in chat_history:
+            for msg in limited_history:
                 filtered_msg = chat_utils.message_copy(msg)
                 if hasattr(filtered_msg, "blocks") and filtered_msg.blocks:
                     filtered_msg.blocks = [
@@ -379,9 +409,10 @@ class CodeActAgent(Workflow):
                 chat_history=chat_history_str,
                 response=response_str,
                 timestamp=time.time(),
-                screenshot=(await ctx.get("screenshot", None))
+                screenshot=(await ctx.store.get("screenshot", None))
             )
 
+
             self.episodic_memory.steps.append(step)
 
             assert hasattr(
@@ -399,41 +430,74 @@ class CodeActAgent(Workflow):
                     logger.error(f"Rate limit error. Retrying in {seconds} seconds...")
                     time.sleep(seconds)
                 else:
-                    logger.error(f"Rate limit error. Retrying in 5 seconds...")
+                    logger.error("Rate limit error. Retrying in 5 seconds...")
                     time.sleep(40)
                 logger.debug("🔍 Retrying call to LLM...")
                 response = await self.llm.achat(messages=messages_to_send)
+            elif (
+                self.llm.class_name() == "Anthropic_LLM"
+                and "overloaded_error" in str(e)
+            ):
+                # Use exponential backoff for Anthropic errors
+                if not hasattr(self, '_anthropic_retry_count'):
+                    self._anthropic_retry_count = 0
+                self._anthropic_retry_count += 1
+                seconds = min(2 ** self._anthropic_retry_count, 60)  # Cap at 60 seconds
+                logger.error(f"Anthropic overload error. Retrying in {seconds} seconds... (attempt {self._anthropic_retry_count})")
+                time.sleep(seconds)
+                logger.debug("🔍 Retrying call to LLM...")
+                response = await self.llm.achat(messages=messages_to_send)
+                self._anthropic_retry_count = 0  # Reset on success
             else:
                 logger.error(f"Could not get an answer from LLM: {repr(e)}")
                 raise e
         logger.debug("  - Received response from LLM.")
         return response
 
+    def _limit_history(
+        self, chat_history: List[ChatMessage]
+    ) -> List[ChatMessage]:
+        if LLM_HISTORY_LIMIT <= 0:
+            return chat_history
+
+        max_messages = LLM_HISTORY_LIMIT * 2
+        if len(chat_history) <= max_messages:
+            return chat_history
+
+        preserved_head: List[ChatMessage] = []
+        if chat_history and chat_history[0].role == "user":
+            preserved_head = [chat_history[0]]
+
+        tail = chat_history[-max_messages:]
+        if preserved_head and preserved_head[0] in tail:
+            preserved_head = []
+
+        return preserved_head + tail
+
     async def _add_final_state_observation(self, ctx: Context) -> None:
         """Add the current UI state and screenshot as the final observation step."""
         try:
             # Get current screenshot and UI state
             screenshot = None
-            ui_state = None
-            
+
             try:
                 _, screenshot_bytes = self.tools.take_screenshot()
                 screenshot = screenshot_bytes
             except Exception as e:
                 logger.warning(f"Failed to capture final screenshot: {e}")
-            
+
             try:
                 (a11y_tree, phone_state) = self.tools.get_state()
             except Exception as e:
                 logger.warning(f"Failed to capture final UI state: {e}")
-            
+
             # Create final observation chat history and response
             final_chat_history = [{"role": "system", "content": "Final state observation after task completion"}]
             final_response = {
-                "role": "user", 
+                "role": "user",
                 "content": f"Final State Observation:\nUI State: {a11y_tree}\nScreenshot: {'Available' if screenshot else 'Not available'}"
             }
-            
+
             # Create final episodic memory step
             final_step = EpisodicMemoryStep(
                 chat_history=json.dumps(final_chat_history),
@@ -441,9 +505,9 @@ class CodeActAgent(Workflow):
                 timestamp=time.time(),
                 screenshot=screenshot
             )
-            
+
             self.episodic_memory.steps.append(final_step)
             logger.info("Added final state observation to episodic memory")
-            
+
         except Exception as e:
             logger.error(f"Failed to add final state observation: {e}")

droidrun/agent/codeact/events.py

@@ -1,10 +1,13 @@
+from typing import Optional
+
 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):
     input: list[ChatMessage]
 
@@ -12,7 +15,7 @@ class TaskInputEvent(Event):
 
 class TaskThinkingEvent(Event):
     thoughts: Optional[str] = None
-    code: Optional[str] = None  
+    code: Optional[str] = None
     usage: Optional[UsageResult] = None
 
 class TaskExecutionEvent(Event):
@@ -28,4 +31,4 @@ class TaskEndEvent(Event):
     reason: str
 
 class EpisodicMemoryEvent(Event):
-    episodic_memory: EpisodicMemory
+    episodic_memory: EpisodicMemory

droidrun/agent/codeact/prompts.py

@@ -21,6 +21,6 @@ Now, describe the next step you will take to address the original goal: {goal}""
 
 # Export all prompts
 __all__ = [
-    "DEFAULT_CODE_ACT_USER_PROMPT", 
+    "DEFAULT_CODE_ACT_USER_PROMPT",
     "DEFAULT_NO_THOUGHTS_PROMPT"
-] 
+]

droidrun/agent/common/constants.py

@@ -0,0 +1,2 @@
+"""Max number of recent conversation steps to include in LLM prompt"""
+LLM_HISTORY_LIMIT = 20

droidrun/agent/common/events.py

@@ -1,5 +1,7 @@
+from typing import Any, Dict
+
 from llama_index.core.workflow import Event
-from typing import Dict, Any
+
 
 class ScreenshotEvent(Event):
     screenshot: bytes
@@ -16,7 +18,7 @@ class TapActionEvent(MacroEvent):
     element_index: int = None
     element_text: str = ""
     element_bounds: str = ""
-    
+
 class SwipeActionEvent(MacroEvent):
     """Event for swipe actions with coordinates"""
     start_x: int
@@ -48,4 +50,4 @@ class StartAppEvent(MacroEvent):
     activity: str = None
 
 class RecordUIStateEvent(Event):
-    ui_state: list[Dict[str, Any]]
+    ui_state: list[Dict[str, Any]]

droidrun/agent/context/__init__.py

@@ -9,15 +9,13 @@ This module contains:
 from .agent_persona import AgentPersona
 from .context_injection_manager import ContextInjectionManager
 from .episodic_memory import EpisodicMemory, EpisodicMemoryStep
-from .reflection import Reflection
-from .task_manager import TaskManager, Task
+from .task_manager import Task, TaskManager
 
 __all__ = [
     "AgentPersona",
     "ContextInjectionManager",
     "EpisodicMemory",
     "EpisodicMemoryStep",
-    "Reflection",
     "TaskManager",
     "Task"
 ]

droidrun/agent/context/agent_persona.py

@@ -1,5 +1,6 @@
-from typing import Dict, List, Callable, Any, Optional
 from dataclasses import dataclass
+from typing import List
+
 
 @dataclass
 class AgentPersona:

droidrun/agent/context/context_injection_manager.py

@@ -5,12 +5,12 @@ This module provides the ContextInjectionManager class that manages different ag
 each with specific system prompts, contexts, and tool subsets tailored for specialized tasks.
 """
 
-import logging
-from typing import Optional, List
-from droidrun.agent.context.agent_persona import AgentPersona
 #import chromadb
 import json
-from pathlib import Path
+import logging
+from typing import List, Optional
+
+from droidrun.agent.context.agent_persona import AgentPersona
 
 logger = logging.getLogger("droidrun")
 
@@ -59,8 +59,8 @@ class ContextInjectionManager:
         Returns:
             AgentPersona instance or None if not found
         """
-        
+
         return self.personas.get(agent_type)
-        
+
     def get_all_personas(self) -> List[str]:
         return self.personas

droidrun/agent/context/episodic_memory.py

@@ -1,7 +1,9 @@
 from dataclasses import dataclass, field
-from droidrun.agent.context.agent_persona import AgentPersona
 from typing import List, Optional
 
+from droidrun.agent.context.agent_persona import AgentPersona
+
+
 @dataclass
 class EpisodicMemoryStep:
     chat_history: str
@@ -9,7 +11,7 @@ class EpisodicMemoryStep:
     timestamp: float
     screenshot: Optional[bytes]
 
-@dataclass 
+@dataclass
 class EpisodicMemory:
     persona: AgentPersona
-    steps: List[EpisodicMemoryStep] = field(default_factory=list)
+    steps: List[EpisodicMemoryStep] = field(default_factory=list)

droidrun/agent/context/personas/__init__.py

@@ -1,11 +1,11 @@
-from .default import DEFAULT
-from .ui_expert import UI_EXPERT
 from .app_starter import APP_STARTER_EXPERT
 from .big_agent import BIG_AGENT
+from .default import DEFAULT
+from .ui_expert import UI_EXPERT
 
 __all__ = [
     'DEFAULT',
     'UI_EXPERT',
     'APP_STARTER_EXPERT',
     'BIG_AGENT',
-    ]
+    ]

droidrun/agent/context/personas/app_starter.py

@@ -2,7 +2,7 @@ from droidrun.agent.context.agent_persona import AgentPersona
 from droidrun.tools import Tools
 
 APP_STARTER_EXPERT = AgentPersona(
-    name="AppStarterExpert", 
+    name="AppStarterExpert",
     description="Specialized in app launching",
     expertise_areas=[
         "app launching"
@@ -37,8 +37,8 @@ APP_STARTER_EXPERT = AgentPersona(
     In addition to the Python Standard Library and any functions you have already written, you can use the following functions:
     {tool_descriptions}
 
-    Reminder: Always place your Python code between ```...``` tags when you want to run code. 
+    Reminder: Always place your Python code between ```...``` tags when you want to run code.
 
     You focus ONLY on app launching and package management - UI interactions within apps are handled by UI specialists.""",
 
-)
+)

droidrun/agent/context/personas/big_agent.py

@@ -5,7 +5,7 @@ BIG_AGENT = AgentPersona(
     name="Big Agent",
     description="Big Agent. Use this as your Big Agent",
     expertise_areas=[
-        "UI navigation", "button interactions", "text input", 
+        "UI navigation", "button interactions", "text input",
         "menu navigation", "form filling", "scrolling", "app launching"
     ],
     allowed_tools=[
@@ -90,7 +90,7 @@ BIG_AGENT = AgentPersona(
     - Present the results clearly and concisely as if you computed them directly
     - Structure your response like you're directly answering the user's query, not explaining how you solved it
 
-    Reminder: Always place your Python code between ```...``` tags when you want to run code. 
+    Reminder: Always place your Python code between ```...``` tags when you want to run code.
 """
 
-)
+)

droidrun/agent/context/personas/default.py

@@ -5,7 +5,7 @@ DEFAULT = AgentPersona(
     name="Default",
     description="Default Agent. Use this as your Default",
     expertise_areas=[
-        "UI navigation", "button interactions", "text input", 
+        "UI navigation", "button interactions", "text input",
         "menu navigation", "form filling", "scrolling", "app launching"
     ],
     allowed_tools=[
@@ -89,7 +89,7 @@ DEFAULT = AgentPersona(
     - Present the results clearly and concisely as if you computed them directly
     - Structure your response like you're directly answering the user's query, not explaining how you solved it
 
-    Reminder: Always place your Python code between ```...``` tags when you want to run code. 
+    Reminder: Always place your Python code between ```...``` tags when you want to run code.
 """
 
-)
+)

droidrun/agent/context/personas/ui_expert.py

@@ -5,14 +5,14 @@ UI_EXPERT = AgentPersona(
     name="UIExpert",
     description="Specialized in UI interactions, navigation, and form filling",
     expertise_areas=[
-        "UI navigation", "button interactions", "text input", 
+        "UI navigation", "button interactions", "text input",
         "menu navigation", "form filling", "scrolling"
     ],
     allowed_tools=[
         Tools.swipe.__name__,
         Tools.input_text.__name__,
         Tools.press_key.__name__,
-        Tools.tap_by_index.__name__, 
+        Tools.tap_by_index.__name__,
         Tools.drag.__name__,
         Tools.remember.__name__,
         Tools.complete.__name__
@@ -54,8 +54,8 @@ UI_EXPERT = AgentPersona(
     - Remember important UI state information for context
 
     You do NOT handle app launching or package management - that's handled by other specialists.
-    
-    
+
+
     ## Available Context:
     In your execution environment, you have access to:
     - `ui_elements`: A global variable containing the current UI elements from the device. This is automatically updated before each code execution and contains the latest UI elements that were fetched.
@@ -99,7 +99,7 @@ UI_EXPERT = AgentPersona(
     - If relevant, you can briefly mention general methods used, but don't include code snippets in the final answer
     - Structure your response like you're directly answering the user's query, not explaining how you solved it
 
-    Reminder: Always place your Python code between ```...``` tags when you want to run code. 
+    Reminder: Always place your Python code between ```...``` tags when you want to run code.
 
     You MUST ALWAYS to include your reasoning and thought process outside of the code block. You MUST DOUBLE CHECK that TASK IS COMPLETE with a SCREENSHOT.
     """