baqr/computer_use_ootb

1import subprocess2import base643from pathlib import Path4from PIL import ImageGrab5from uuid import uuid46from screeninfo import get_monitors7import platform8if platform.system() == "Darwin":9    import Quartz  # uncomment this line if you are on macOS10    11from PIL import ImageGrab12from functools import partial13from .base import BaseAnthropicTool, ToolError, ToolResult14 15 16OUTPUT_DIR = "./tmp/outputs"17 18def get_screenshot(selected_screen: int = 0, resize: bool = True, target_width: int = 1920, target_height: int = 1080):19        # print(f"get_screenshot selected_screen: {selected_screen}")20        21        # Get screen width and height using Windows command22        display_num = None23        offset_x = 024        offset_y = 025        selected_screen = selected_screen   26        width, height = _get_screen_size()    27    28        """Take a screenshot of the current screen and return a ToolResult with the base64 encoded image."""29        output_dir = Path(OUTPUT_DIR)30        output_dir.mkdir(parents=True, exist_ok=True)31        path = output_dir / f"screenshot_{uuid4().hex}.png"32 33        ImageGrab.grab = partial(ImageGrab.grab, all_screens=True)34 35        # Detect platform36        system = platform.system()37 38        if system == "Windows":39            # Windows: Use screeninfo to get monitor details40            screens = get_monitors()41 42            # Sort screens by x position to arrange from left to right43            sorted_screens = sorted(screens, key=lambda s: s.x)44 45            if selected_screen < 0 or selected_screen >= len(screens):46                raise IndexError("Invalid screen index.")47 48            screen = sorted_screens[selected_screen]49            bbox = (screen.x, screen.y, screen.x + screen.width, screen.y + screen.height)50 51        elif system == "Darwin":  # macOS52            # macOS: Use Quartz to get monitor details53            max_displays = 32  # Maximum number of displays to handle54            active_displays = Quartz.CGGetActiveDisplayList(max_displays, None, None)[1]55 56            # Get the display bounds (resolution) for each active display57            screens = []58            for display_id in active_displays:59                bounds = Quartz.CGDisplayBounds(display_id)60                screens.append({61                    'id': display_id,62                    'x': int(bounds.origin.x),63                    'y': int(bounds.origin.y),64                    'width': int(bounds.size.width),65                    'height': int(bounds.size.height),66                    'is_primary': Quartz.CGDisplayIsMain(display_id)  # Check if this is the primary display67                })68 69            # Sort screens by x position to arrange from left to right70            sorted_screens = sorted(screens, key=lambda s: s['x'])71            # print(f"Darwin sorted_screens: {sorted_screens}")72 73            if selected_screen < 0 or selected_screen >= len(screens):74                raise IndexError("Invalid screen index.")75 76            screen = sorted_screens[selected_screen]77            78            bbox = (screen['x'], screen['y'], screen['x'] + screen['width'], screen['y'] + screen['height'])79 80        else:  # Linux or other OS81            cmd = "xrandr | grep ' primary' | awk '{print $4}'"82            try:83                output = subprocess.check_output(cmd, shell=True).decode()84                resolution = output.strip().split()[0]85                width, height = map(int, resolution.split('x'))86                bbox = (0, 0, width, height)  # Assuming single primary screen for simplicity87            except subprocess.CalledProcessError:88                raise RuntimeError("Failed to get screen resolution on Linux.")89 90        # Take screenshot using the bounding box91        screenshot = ImageGrab.grab(bbox=bbox)92 93        # Set offsets (for potential future use)94        offset_x = screen['x'] if system == "Darwin" else screen.x95        offset_y = screen['y'] if system == "Darwin" else screen.y96 97        # # Resize if 98        if resize:99            screenshot = screenshot.resize((target_width, target_height))100 101        # Save the screenshot102        screenshot.save(str(path))103 104        if path.exists():105            # Return a ToolResult instance instead of a dictionary106            return screenshot, path107        108        raise ToolError(f"Failed to take screenshot: {path} does not exist.")109    110    111 112 113def _get_screen_size(selected_screen: int = 0):114    if platform.system() == "Windows":115        # Use screeninfo to get primary monitor on Windows116        screens = get_monitors()117 118        # Sort screens by x position to arrange from left to right119        sorted_screens = sorted(screens, key=lambda s: s.x)120        if selected_screen is None:121            primary_monitor = next((m for m in get_monitors() if m.is_primary), None)122            return primary_monitor.width, primary_monitor.height123        elif selected_screen < 0 or selected_screen >= len(screens):124            raise IndexError("Invalid screen index.")125        else:126            screen = sorted_screens[selected_screen]127            return screen.width, screen.height128    elif platform.system() == "Darwin":129        # macOS part using Quartz to get screen information130        max_displays = 32  # Maximum number of displays to handle131        active_displays = Quartz.CGGetActiveDisplayList(max_displays, None, None)[1]132 133        # Get the display bounds (resolution) for each active display134        screens = []135        for display_id in active_displays:136            bounds = Quartz.CGDisplayBounds(display_id)137            screens.append({138                'id': display_id,139                'x': int(bounds.origin.x),140                'y': int(bounds.origin.y),141                'width': int(bounds.size.width),142                'height': int(bounds.size.height),143                'is_primary': Quartz.CGDisplayIsMain(display_id)  # Check if this is the primary display144            })145 146        # Sort screens by x position to arrange from left to right147        sorted_screens = sorted(screens, key=lambda s: s['x'])148 149        if selected_screen is None:150            # Find the primary monitor151            primary_monitor = next((screen for screen in screens if screen['is_primary']), None)152            if primary_monitor:153                return primary_monitor['width'], primary_monitor['height']154            else:155                raise RuntimeError("No primary monitor found.")156        elif selected_screen < 0 or selected_screen >= len(screens):157            raise IndexError("Invalid screen index.")158        else:159            # Return the resolution of the selected screen160            screen = sorted_screens[selected_screen]161            return screen['width'], screen['height']162 163    else:  # Linux or other OS164        cmd = "xrandr | grep ' primary' | awk '{print $4}'"165        try:166            output = subprocess.check_output(cmd, shell=True).decode()167            resolution = output.strip().split()[0]168            width, height = map(int, resolution.split('x'))169            return width, height170        except subprocess.CalledProcessError:171            raise RuntimeError("Failed to get screen resolution on Linux.")172