droidrun 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

droidrun/agent/utils/trajectory.py

@@ -18,10 +18,24 @@ logger = logging.getLogger("droidrun")
 
 class Trajectory:
 
-    def __init__(self):
-        """Initializes an empty trajectory class."""
+    def __init__(self, goal: str = None):
+        """Initializes an empty trajectory class.
+        
+        Args:
+            goal: The goal/prompt that this trajectory is trying to achieve
+        """
         self.events: List[Event] = [] 
         self.screenshots: List[bytes] = [] 
+        self.macro: List[Event] = []
+        self.goal = goal or "DroidRun automation sequence"
+
+    def set_goal(self, goal: str) -> None:
+        """Update the goal/description for this trajectory.
+        
+        Args:
+            goal: The new goal/prompt description
+        """
+        self.goal = goal
 
 
     def create_screenshot_gif(self, output_path: str, duration: int = 1000) -> str:
@@ -33,9 +47,10 @@ class Trajectory:
             duration: Duration for each frame in milliseconds
         
         Returns:
-            Path to the created GIF file
+            Path to the created GIF file, or None if no screenshots available
         """
         if len(self.screenshots) == 0:
+            logger.info("📷 No screenshots available for GIF creation")
             return None
             
         images = []
@@ -62,17 +77,20 @@ class Trajectory:
     ) -> str:
         """
         Save trajectory steps to a JSON file and create a GIF of screenshots if available.
+        Also saves the macro sequence as a separate file for replay.
+        Creates a dedicated folder for each trajectory containing all related files.
         
         Args:
-            directory: Directory to save the trajectory files
+            directory: Base directory to save the trajectory files
         
         Returns:
-            Path to the saved trajectory file
+            Path to the trajectory folder
         """
         os.makedirs(directory, exist_ok=True)
         
         timestamp = time.strftime("%Y%m%d_%H%M%S")
-        base_path = os.path.join(directory, f"trajectory_{timestamp}")
+        trajectory_folder = os.path.join(directory, f"trajectory_{timestamp}")
+        os.makedirs(trajectory_folder, exist_ok=True)
         
         def make_serializable(obj):
             """Recursively make objects JSON serializable."""
@@ -100,6 +118,7 @@ class Trajectory:
             else:
                 return obj
         
+        # Save main trajectory events
         serializable_events = []
         for event in self.events:
             event_dict = {
@@ -109,13 +128,209 @@ class Trajectory:
             }
             serializable_events.append(event_dict)
         
-        json_path = f"{base_path}.json"
-        with open(json_path, "w") as f:
+        trajectory_json_path = os.path.join(trajectory_folder, "trajectory.json")
+        with open(trajectory_json_path, "w") as f:
             json.dump(serializable_events, f, indent=2)
 
-        self.create_screenshot_gif(base_path)
+        # Save macro sequence as a separate file for replay
+        if self.macro:
+            macro_data = []
+            for macro_event in self.macro:
+                macro_dict = {
+                    "type": macro_event.__class__.__name__,
+                    **{k: make_serializable(v) for k, v in macro_event.__dict__.items() 
+                       if not k.startswith('_')}
+                }
+                macro_data.append(macro_dict)
+            
+            macro_json_path = os.path.join(trajectory_folder, "macro.json")
+            with open(macro_json_path, "w") as f:
+                json.dump({
+                    "version": "1.0",
+                    "description": self.goal,
+                    "timestamp": timestamp,
+                    "total_actions": len(macro_data),
+                    "actions": macro_data
+                }, f, indent=2)
+            
+            logger.info(f"💾 Saved macro sequence with {len(macro_data)} actions to {macro_json_path}")
+
+        # Create screenshot GIF
+        gif_path = self.create_screenshot_gif(os.path.join(trajectory_folder, "screenshots"))
+        if gif_path:
+            logger.info(f"🎬 Saved screenshot GIF to {gif_path}")
+
+        logger.info(f"📁 Trajectory saved to folder: {trajectory_folder}")
+        return trajectory_folder
+
+    @staticmethod
+    def load_trajectory_folder(trajectory_folder: str) -> Dict[str, Any]:
+        """
+        Load trajectory data from a trajectory folder.
+        
+        Args:
+            trajectory_folder: Path to the trajectory folder
+            
+        Returns:
+            Dictionary containing trajectory data, macro data, and file paths
+        """
+        result = {
+            "trajectory_data": None,
+            "macro_data": None,
+            "gif_path": None,
+            "folder_path": trajectory_folder
+        }
+        
+        try:
+            # Load main trajectory
+            trajectory_json_path = os.path.join(trajectory_folder, "trajectory.json")
+            if os.path.exists(trajectory_json_path):
+                with open(trajectory_json_path, "r") as f:
+                    result["trajectory_data"] = json.load(f)
+                logger.info(f"📖 Loaded trajectory data from {trajectory_json_path}")
+            
+            # Load macro sequence
+            macro_json_path = os.path.join(trajectory_folder, "macro.json")
+            if os.path.exists(macro_json_path):
+                with open(macro_json_path, "r") as f:
+                    result["macro_data"] = json.load(f)
+                logger.info(f"📖 Loaded macro data from {macro_json_path}")
+            
+            # Check for GIF
+            gif_path = os.path.join(trajectory_folder, "screenshots.gif")
+            if os.path.exists(gif_path):
+                result["gif_path"] = gif_path
+                logger.info(f"🎬 Found screenshot GIF at {gif_path}")
+            
+            return result
+            
+        except Exception as e:
+            logger.error(f"❌ Error loading trajectory folder {trajectory_folder}: {e}")
+            return result
+
+    @staticmethod
+    def load_macro_sequence(macro_file_path: str) -> Dict[str, Any]:
+        """
+        Load a macro sequence from a saved macro file.
+        
+        Args:
+            macro_file_path: Path to the macro JSON file (can be full path or trajectory folder)
+            
+        Returns:
+            Dictionary containing the macro sequence data
+        """
+        # Check if it's a folder path - if so, look for macro.json inside
+        if os.path.isdir(macro_file_path):
+            macro_file_path = os.path.join(macro_file_path, "macro.json")
+        
+        try:
+            with open(macro_file_path, "r") as f:
+                macro_data = json.load(f)
+            
+            logger.info(f"📖 Loaded macro sequence with {macro_data.get('total_actions', 0)} actions from {macro_file_path}")
+            return macro_data
+        except FileNotFoundError:
+            logger.error(f"❌ Macro file not found: {macro_file_path}")
+            return {}
+        except json.JSONDecodeError as e:
+            logger.error(f"❌ Error parsing macro file {macro_file_path}: {e}")
+            return {}
+
+    @staticmethod
+    def get_macro_summary(macro_data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Get a summary of a macro sequence.
+        
+        Args:
+            macro_data: The macro data dictionary
+            
+        Returns:
+            Dictionary with statistics about the macro
+        """
+        if not macro_data or "actions" not in macro_data:
+            return {"error": "Invalid macro data"}
+        
+        actions = macro_data["actions"]
+        
+        # Count action types
+        action_types = {}
+        for action in actions:
+            action_type = action.get("action_type", "unknown")
+            action_types[action_type] = action_types.get(action_type, 0) + 1
+        
+        # Calculate duration if timestamps are available
+        timestamps = [action.get("timestamp") for action in actions if action.get("timestamp")]
+        duration = max(timestamps) - min(timestamps) if len(timestamps) > 1 else 0
+        
+        return {
+            "version": macro_data.get("version", "unknown"),
+            "description": macro_data.get("description", "No description"),
+            "total_actions": len(actions),
+            "action_types": action_types,
+            "duration_seconds": round(duration, 2) if duration > 0 else None,
+            "timestamp": macro_data.get("timestamp", "unknown")
+        }
 
-        return json_path
+    @staticmethod
+    def print_macro_summary(macro_file_path: str) -> None:
+        """
+        Print a summary of a macro sequence.
+        
+        Args:
+            macro_file_path: Path to the macro JSON file or trajectory folder
+        """
+        macro_data = Trajectory.load_macro_sequence(macro_file_path)
+        if not macro_data:
+            print("❌ Could not load macro data")
+            return
+            
+        summary = Trajectory.get_macro_summary(macro_data)
+        
+        print("=== Macro Summary ===")
+        print(f"File: {macro_file_path}")
+        print(f"Version: {summary.get('version', 'unknown')}")
+        print(f"Description: {summary.get('description', 'No description')}")
+        print(f"Timestamp: {summary.get('timestamp', 'unknown')}")
+        print(f"Total actions: {summary.get('total_actions', 0)}")
+        if summary.get('duration_seconds'):
+            print(f"Duration: {summary['duration_seconds']} seconds")
+        print("Action breakdown:")
+        for action_type, count in summary.get('action_types', {}).items():
+            print(f"  - {action_type}: {count}")
+        print("=====================")
+
+    @staticmethod
+    def print_trajectory_folder_summary(trajectory_folder: str) -> None:
+        """
+        Print a comprehensive summary of a trajectory folder.
+        
+        Args:
+            trajectory_folder: Path to the trajectory folder
+        """
+        folder_data = Trajectory.load_trajectory_folder(trajectory_folder)
+        
+        print("=== Trajectory Folder Summary ===")
+        print(f"Folder: {trajectory_folder}")
+        print(f"Trajectory data: {'✅ Available' if folder_data['trajectory_data'] else '❌ Missing'}")
+        print(f"Macro data: {'✅ Available' if folder_data['macro_data'] else '❌ Missing'}")
+        print(f"Screenshot GIF: {'✅ Available' if folder_data['gif_path'] else '❌ Missing'}")
+        
+        if folder_data['macro_data']:
+            print("\n--- Macro Summary ---")
+            summary = Trajectory.get_macro_summary(folder_data['macro_data'])
+            print(f"Description: {summary.get('description', 'No description')}")
+            print(f"Total actions: {summary.get('total_actions', 0)}")
+            if summary.get('duration_seconds'):
+                print(f"Duration: {summary['duration_seconds']} seconds")
+            print("Action breakdown:")
+            for action_type, count in summary.get('action_types', {}).items():
+                print(f"  - {action_type}: {count}")
+        
+        if folder_data['trajectory_data']:
+            print(f"\n--- Trajectory Summary ---")
+            print(f"Total events: {len(folder_data['trajectory_data'])}")
+        
+        print("=================================")
 
     def get_trajectory_statistics(trajectory_data: Dict[str, Any]) -> Dict[str, Any]:
         """
@@ -181,4 +396,36 @@ class Trajectory:
         print(f"Execution steps: {stats['execution_steps']}")
         print(f"Successful executions: {stats['successful_executions']}")
         print(f"Failed executions: {stats['failed_executions']}")
-        print("==========================")
+        print("==========================")
+
+
+# Example usage:
+"""
+# Save a trajectory with a specific goal (automatically creates folder structure)
+trajectory = Trajectory(goal="Open settings and check battery level")
+# ... add events and screenshots to trajectory ...
+folder_path = trajectory.save_trajectory()
+
+# Or update the goal later
+trajectory.set_goal("Navigate to Settings and find device info")
+
+# Load entire trajectory folder
+folder_data = Trajectory.load_trajectory_folder(folder_path)
+trajectory_events = folder_data['trajectory_data']
+macro_actions = folder_data['macro_data']
+gif_path = folder_data['gif_path']
+
+# Load just the macro from folder
+macro_data = Trajectory.load_macro_sequence(folder_path)
+
+# Print summaries
+Trajectory.print_trajectory_folder_summary(folder_path)
+Trajectory.print_macro_summary(folder_path)
+
+# Example folder structure created:
+# trajectories/
+# └── trajectory_20250108_143052/
+#     ├── trajectory.json      # Full trajectory events
+#     ├── macro.json          # Macro sequence with goal as description
+#     └── screenshots.gif     # Screenshot animation
+"""

droidrun/cli/main.py

@@ -9,10 +9,11 @@ import logging
 import warnings
 from contextlib import nullcontext
 from rich.console import Console
+from adbutils import adb
 from droidrun.agent.droid import DroidAgent
 from droidrun.agent.utils.llm_picker import load_llm
-from droidrun.adb import DeviceManager
 from droidrun.tools import AdbTools, IOSTools
+from droidrun.agent.context.personas import DEFAULT, BIG_AGENT
 from functools import wraps
 from droidrun.cli.logs import LogHandler
 from droidrun.telemetry import print_telemetry_message
@@ -21,7 +22,10 @@ from droidrun.portal import (
     enable_portal_accessibility,
     PORTAL_PACKAGE_NAME,
     ping_portal,
+    ping_portal_tcp,
+    ping_portal_content,
 )
+from droidrun.macro.cli import macro_cli
 
 # Suppress all warnings
 warnings.filterwarnings("ignore")
@@ -29,7 +33,6 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false"
 
 console = Console()
-device_manager = DeviceManager()
 
 
 def configure_logging(goal: str, debug: bool):
@@ -38,7 +41,7 @@ def configure_logging(goal: str, debug: bool):
 
     handler = LogHandler(goal)
     handler.setFormatter(
-        logging.Formatter("%(levelname)s %(message)s", "%H:%M:%S")
+        logging.Formatter("%(levelname)s %(name)s %(message)s", "%H:%M:%S")
         if debug
         else logging.Formatter("%(message)s", "%H:%M:%S")
     )
@@ -47,6 +50,12 @@ def configure_logging(goal: str, debug: bool):
     logger.setLevel(logging.DEBUG if debug else logging.INFO)
     logger.propagate = False
 
+    if debug:
+        tools_logger = logging.getLogger("droidrun-tools")
+        tools_logger.addHandler(handler)
+        tools_logger.propagate = False
+        tools_logger.setLevel(logging.DEBUG if debug else logging.INFO)
+
     return handler
 
 
@@ -72,8 +81,10 @@ async def run_command(
     reflection: bool,
     tracing: bool,
     debug: bool,
-    save_trajectory: bool = False,
+    use_tcp: bool,
+    save_trajectory: str = "none",
     ios: bool = False,
+    allow_drag: bool = False,
     **kwargs,
 ):
     """Run a command on your Android device using natural language."""
@@ -95,8 +106,8 @@ async def run_command(
             # Device setup
             if device is None and not ios:
                 logger.info("🔍 Finding connected device...")
-                device_manager = DeviceManager()
-                devices = await device_manager.list_devices()
+
+                devices = adb.list()
                 if not devices:
                     raise ValueError("No connected devices found.")
                 device = devices[0].serial
@@ -108,7 +119,16 @@ async def run_command(
             else:
                 logger.info(f"📱 Using device: {device}")
 
-            tools = AdbTools(serial=device) if not ios else IOSTools(url=device)
+            tools = (
+                AdbTools(serial=device, use_tcp=use_tcp)
+                if not ios
+                else IOSTools(url=device)
+            )
+            # Set excluded tools based on CLI flags
+            excluded_tools = [] if allow_drag else ["drag"]
+
+            # Select personas based on --drag flag
+            personas = [BIG_AGENT] if allow_drag else [DEFAULT]
 
             # LLM setup
             log_handler.update_step("Initializing LLM...")
@@ -134,6 +154,8 @@ async def run_command(
                 goal=command,
                 llm=llm,
                 tools=tools,
+                personas=personas,
+                excluded_tools=excluded_tools,
                 max_steps=steps,
                 timeout=1000,
                 vision=vision,
@@ -237,11 +259,17 @@ class DroidRunCLI(click.Group):
     "--debug", is_flag=True, help="Enable verbose debug logging", default=False
 )
 @click.option(
-    "--save-trajectory",
+    "--use-tcp",
     is_flag=True,
-    help="Save agent trajectory to file",
+    help="Use TCP communication for device control",
     default=False,
 )
+@click.option(
+    "--save-trajectory",
+    type=click.Choice(["none", "step", "action"]),
+    help="Trajectory saving level: none (no saving), step (save per step), action (save per action)",
+    default="none",
+)
 @click.group(cls=DroidRunCLI)
 def cli(
     device: str | None,
@@ -256,7 +284,8 @@ def cli(
     reflection: bool,
     tracing: bool,
     debug: bool,
-    save_trajectory: bool,
+    use_tcp: bool,
+    save_trajectory: str,
 ):
     """DroidRun - Control your Android device through LLM agents."""
     pass
@@ -311,10 +340,23 @@ def cli(
 @click.option(
     "--debug", is_flag=True, help="Enable verbose debug logging", default=False
 )
+@click.option(
+    "--use-tcp",
+    is_flag=True,
+    help="Use TCP communication for device control",
+    default=False,
+)
 @click.option(
     "--save-trajectory",
+    type=click.Choice(["none", "step", "action"]),
+    help="Trajectory saving level: none (no saving), step (save per step), action (save per action)",
+    default="none",
+)
+@click.option(
+    "--drag",
+    "allow_drag",
     is_flag=True,
-    help="Save agent trajectory to file",
+    help="Enable drag tool",
     default=False,
 )
 @click.option("--ios", is_flag=True, help="Run on iOS device", default=False)
@@ -332,7 +374,9 @@ def run(
     reflection: bool,
     tracing: bool,
     debug: bool,
-    save_trajectory: bool,
+    use_tcp: bool,
+    save_trajectory: str,
+    allow_drag: bool,
     ios: bool,
 ):
     """Run a command on your Android device using natural language."""
@@ -350,18 +394,19 @@ def run(
         reflection,
         tracing,
         debug,
+        use_tcp,
         temperature=temperature,
         save_trajectory=save_trajectory,
+        allow_drag=allow_drag,
         ios=ios,
     )
 
 
 @cli.command()
-@coro
-async def devices():
+def devices():
     """List connected Android devices."""
     try:
-        devices = await device_manager.list_devices()
+        devices = adb.list()
         if not devices:
             console.print("[yellow]No devices connected.[/]")
             return
@@ -375,27 +420,24 @@ async def devices():
 
 @cli.command()
 @click.argument("serial")
-@click.option("--port", "-p", default=5555, help="ADB port (default: 5555)")
-@coro
-async def connect(serial: str, port: int):
+def connect(serial: str):
     """Connect to a device over TCP/IP."""
     try:
-        device = await device_manager.connect(serial, port)
-        if device:
-            console.print(f"[green]Successfully connected to {serial}:{port}[/]")
+        device = adb.connect(serial)
+        if device.count("already connected"):
+            console.print(f"[green]Successfully connected to {serial}[/]")
         else:
-            console.print(f"[red]Failed to connect to {serial}:{port}[/]")
+            console.print(f"[red]Failed to connect to {serial}: {device}[/]")
     except Exception as e:
         console.print(f"[red]Error connecting to device: {e}[/]")
 
 
 @cli.command()
 @click.argument("serial")
-@coro
-async def disconnect(serial: str):
+def disconnect(serial: str):
     """Disconnect from a device."""
     try:
-        success = await device_manager.disconnect(serial)
+        success = adb.disconnect(serial, raise_error=True)
         if success:
             console.print(f"[green]Successfully disconnected from {serial}[/]")
         else:
@@ -406,16 +448,19 @@ async def disconnect(serial: str):
 
 @cli.command()
 @click.option("--device", "-d", help="Device serial number or IP address", default=None)
-@click.option("--path", help="Path to the Droidrun Portal APK to install on the device. If not provided, the latest portal apk version will be downloaded and installed.", default=None)
+@click.option(
+    "--path",
+    help="Path to the Droidrun Portal APK to install on the device. If not provided, the latest portal apk version will be downloaded and installed.",
+    default=None,
+)
 @click.option(
     "--debug", is_flag=True, help="Enable verbose debug logging", default=False
 )
-@coro
-async def setup(path: str | None, device: str | None, debug: bool):
+def setup(path: str | None, device: str | None, debug: bool):
     """Install and enable the DroidRun Portal on a device."""
     try:
         if not device:
-            devices = await device_manager.list_devices()
+            devices = adb.list()
             if not devices:
                 console.print("[yellow]No devices connected.[/]")
                 return
@@ -423,7 +468,7 @@ async def setup(path: str | None, device: str | None, debug: bool):
             device = devices[0].serial
             console.print(f"[blue]Using device:[/] {device}")
 
-        device_obj = await device_manager.get_device(device)
+        device_obj = adb.device(device)
         if not device_obj:
             console.print(
                 f"[bold red]Error:[/] Could not get device object for {device}"
@@ -443,18 +488,20 @@ async def setup(path: str | None, device: str | None, debug: bool):
                 return
 
             console.print(f"[bold blue]Step 1/2: Installing APK:[/] {apk_path}")
-            result = await device_obj.install_app(apk_path, True, True)
-
-            if "Error" in result:
-                console.print(f"[bold red]Installation failed:[/] {result}")
+            try:
+                device_obj.install(
+                    apk_path, uninstall=True, flags=["-g"], silent=not debug
+                )
+            except Exception as e:
+                console.print(f"[bold red]Installation failed:[/] {e}")
                 return
-            else:
-                console.print(f"[bold green]Installation successful![/]")
+
+            console.print(f"[bold green]Installation successful![/]")
 
             console.print(f"[bold blue]Step 2/2: Enabling accessibility service[/]")
 
             try:
-                await enable_portal_accessibility(device_obj)
+                enable_portal_accessibility(device_obj)
 
                 console.print("[green]Accessibility service enabled successfully![/]")
                 console.print(
@@ -469,9 +516,7 @@ async def setup(path: str | None, device: str | None, debug: bool):
                     "[yellow]Opening accessibility settings for manual configuration...[/]"
                 )
 
-                await device_obj.shell(
-                    "am start -a android.settings.ACCESSIBILITY_SETTINGS"
-                )
+                device_obj.shell("am start -a android.settings.ACCESSIBILITY_SETTINGS")
 
                 console.print(
                     "\n[yellow]Please complete the following steps on your device:[/]"
@@ -500,19 +545,30 @@ async def setup(path: str | None, device: str | None, debug: bool):
 
 @cli.command()
 @click.option("--device", "-d", help="Device serial number or IP address", default=None)
+@click.option(
+    "--use-tcp",
+    is_flag=True,
+    help="Use TCP communication for device control",
+    default=False,
+)
 @click.option(
     "--debug", is_flag=True, help="Enable verbose debug logging", default=False
 )
-@coro
-async def ping(device: str | None, debug: bool):
+def ping(device: str | None, use_tcp: bool, debug: bool):
     """Ping a device to check if it is ready and accessible."""
     try:
-        device_obj = await device_manager.get_device(device)
+        device_obj = adb.device(device)
         if not device_obj:
             console.print(f"[bold red]Error:[/] Could not find device {device}")
             return
 
-        await ping_portal(device_obj, debug)
+        ping_portal(device_obj, debug)
+
+        if use_tcp:
+            ping_portal_tcp(device_obj, debug)
+        else:
+            ping_portal_content(device_obj, debug)
+
         console.print(
             "[bold green]Portal is installed and accessible. You're good to go![/]"
         )
@@ -524,6 +580,10 @@ async def ping(device: str | None, debug: bool):
             traceback.print_exc()
 
 
+# Add macro commands as a subgroup
+cli.add_command(macro_cli, name="macro")
+
+
 if __name__ == "__main__":
     command = "Open the settings app"
     device = None
@@ -537,9 +597,11 @@ if __name__ == "__main__":
     reflection = False
     tracing = True
     debug = True
+    use_tcp = True
     base_url = None
     api_base = None
     ios = False
+    allow_drag = False
     run_command(
         command=command,
         device=device,
@@ -552,8 +614,10 @@ if __name__ == "__main__":
         reflection=reflection,
         tracing=tracing,
         debug=debug,
+        use_tcp=use_tcp,
         base_url=base_url,
         api_base=api_base,
         api_key=api_key,
+        allow_drag=allow_drag,
         ios=ios,
     )