droidrun 0.3.9__py3-none-any.whl → 0.3.10.dev3__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,37 +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.constants import LLM_HISTORY_LIMIT
-from droidrun.agent.common.events import RecordUIStateEvent, ScreenshotEvent
-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")
 
@@ -49,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,
@@ -78,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)
+
+        # Add non-atomic tools (remember, complete) from tools_instance
+        self.tool_list["remember"] = tools_instance.remember
+        self.tool_list["complete"] = tools_instance.complete
 
-        for tool_name in persona.allowed_tools:
-            if tool_name in all_tools_list:
-                self.tool_list[tool_name] = all_tools_list[tool_name]
+        # Get tool descriptions from ATOMIC_ACTION_SIGNATURES and custom_tools
+        self.tool_descriptions = get_atomic_tool_descriptions()
 
-        self.tool_descriptions = chat_utils.parse_tool_descriptions(self.tool_list)
+        # 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
@@ -162,7 +189,7 @@ 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.store.set("remembered_info", self.remembered_info)
             chat_history = await chat_utils.add_memory_block(self.remembered_info, chat_history)
@@ -178,7 +205,7 @@ class CodeActAgent(Workflow):
                     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":
@@ -190,8 +217,8 @@ class CodeActAgent(Workflow):
                         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":
@@ -254,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:
@@ -269,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
@@ -323,11 +350,11 @@ class CodeActAgent(Workflow):
         """Finalize the workflow."""
         self.tools.finished = False
         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(
             {
@@ -403,7 +430,7 @@ 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)
@@ -452,26 +479,25 @@ class CodeActAgent(Workflow):
         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),
@@ -479,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/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

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.
     """

droidrun/agent/context/task_manager.py

@@ -1,7 +1,8 @@
+import copy
 import os
-from typing import List, Dict, Optional
 from dataclasses import dataclass
-import copy
+from typing import Dict, List, Optional
+
 
 @dataclass
 class Task:
@@ -14,13 +15,13 @@ class Task:
     # Optional fields to carry success/failure context back to the planner
     message: Optional[str] = None
     failure_reason: Optional[str] = None
-    
+
 
 class TaskManager:
     """
     Manages a list of tasks for an agent, each with a status and assigned specialized agent.
     """
-    STATUS_PENDING = "pending"      
+    STATUS_PENDING = "pending"
     STATUS_COMPLETED = "completed"
     STATUS_FAILED = "failed"
 
@@ -32,14 +33,14 @@ class TaskManager:
     def __init__(self):
         """Initializes an empty task list."""
         self.tasks: List[Task] = []
-        self.goal_completed = False 
+        self.goal_completed = False
         self.message = None
-        self.task_history = [] 
+        self.task_history = []
         self.file_path = os.path.join(os.path.dirname(__file__), "todo.txt")
 
     def get_all_tasks(self) -> List[Task]:
         return self.tasks
-        
+
     def get_task_history(self):
         return self.task_history
 
@@ -67,9 +68,6 @@ class TaskManager:
 
     def get_failed_tasks(self) -> list[dict]:
         return [task for task in self.task_history if task.status == self.STATUS_FAILED]
-    
-    def get_task_history(self) -> list[dict]:
-        return self.task_history
 
 
     def save_to_file(self):
@@ -89,12 +87,12 @@ class TaskManager:
     def set_tasks_with_agents(self, task_assignments: List[Dict[str, str]]):
         """
         Clears the current task list and sets new tasks with their assigned agents.
-        
+
         Args:
             task_assignments: A list of dictionaries, each containing:
                             - 'task': The task description string
                             - 'agent': The agent type
-                            
+
         Example:
             task_manager.set_tasks_with_agents([
                 {'task': 'Open Gmail app', 'agent': 'AppStarterExpert'},
@@ -106,21 +104,21 @@ class TaskManager:
             for i, assignment in enumerate(task_assignments):
                 if not isinstance(assignment, dict) or 'task' not in assignment:
                     raise ValueError(f"Each task assignment must be a dictionary with 'task' key at index {i}.")
-                
+
                 task_description = assignment['task']
                 if not isinstance(task_description, str) or not task_description.strip():
                     raise ValueError(f"Task description must be a non-empty string at index {i}.")
-                
+
                 agent_type = assignment.get('agent', 'Default')
-                
+
                 task_obj = Task(
                     description=task_description.strip(),
                     status=self.STATUS_PENDING,
                     agent_type=agent_type
                 )
-                
+
                 self.tasks.append(task_obj)
-            
+
             print(f"Tasks set with agents: {len(self.tasks)} tasks added.")
             self.save_to_file()
         except Exception as e:

droidrun/agent/droid/__init__.py

@@ -10,4 +10,4 @@ from droidrun.agent.droid.droid_agent import DroidAgent
 __all__ = [
     "CodeActAgent",
     "DroidAgent"
-] 
+]