wcgw 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

wcgw/__init__.py

@@ -1,2 +1,3 @@
 from .client.cli import app
 from .client.tools import run as listen
+from .client.mcp_server import main as mcp_server

wcgw/client/__main__.py

@@ -1,3 +1,3 @@
-from .tools import run
+from .cli import app
 
-run()
+app()

wcgw/client/anthropic_client.py

@@ -27,15 +27,19 @@ from ..types_ import (
     CreateFileNew,
     FileEditFindReplace,
     FileEdit,
+    Keyboard,
+    Mouse,
     ReadFile,
     ReadImage,
-    Writefile,
     ResetShell,
+    ScreenShot,
+    GetScreenInfo,
 )
 
 from .common import Models, discard_input
 from .common import CostData
 from .tools import ImageData
+from .computer_use import Computer
 
 from .tools import (
     DoneFlag,
@@ -166,6 +170,7 @@ def loop(
 - The first line might be `(...truncated)` if the output is too long.
 - Always run `pwd` if you get any file or directory not found error to make sure you're not lost.
 - The control will return to you in 5 seconds regardless of the status. For heavy commands, keep checking status using BashInteraction till they are finished.
+- Run long running commands in background using screen instead of "&".
 """,
         ),
         ToolParam(
@@ -192,7 +197,6 @@ def loop(
             name="CreateFileNew",
             description="""
 - Write content to a new file. Provide file path and content. Use this instead of BashCommand for writing new files.
-- This doesn't create any directories, please create directories using `mkdir -p` BashCommand.
 - Provide absolute file path only.
 - For editing existing files, use FileEdit instead of this tool.
 """,
@@ -205,7 +209,7 @@ def loop(
         ToolParam(
             input_schema=ResetShell.model_json_schema(),
             name="ResetShell",
-            description="Resets the shell. Use only if all interrupts and prompt reset attempts have failed repeatedly.",
+            description="Resets the shell. Use only if all interrupts and prompt reset attempts have failed repeatedly.\nAlso exits the docker environment.\nYou need to call GetScreenInfo again",
         ),
         ToolParam(
             input_schema=FileEdit.model_json_schema(),
@@ -213,6 +217,46 @@ def loop(
             description="""
 - Use absolute file path only.
 - Use SEARCH/REPLACE blocks to edit the file.
+""",
+        ),
+        ToolParam(
+            input_schema=GetScreenInfo.model_json_schema(),
+            name="GetScreenInfo",
+            description="""
+- Get display information of an OS running on docker using image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
+- If user hasn't provided docker image id, check using `docker ps` and provide the id.
+- Important: call this first in the conversation before ScreenShot, Mouse, and Keyboard tools.
+- Connects shell to the docker environment.
+- Note: once this is called, the shell enters the docker environment. All bash commands will run over there.
+""",
+        ),
+        ToolParam(
+            input_schema=ScreenShot.model_json_schema(),
+            name="ScreenShot",
+            description="""
+- Capture screenshot of an OS running on docker using image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
+- If user hasn't provided docker image id, check using `docker ps` and provide the id.
+- Capture ScreenShot of the current screen for automation.
+""",
+        ),
+        ToolParam(
+            input_schema=Mouse.model_json_schema(),
+            name="Mouse",
+            description="""
+- Interact with docker container running image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
+- If user hasn't provided docker image id, check using `docker ps` and provide the id.
+- Interact with the screen using mouse
+""",
+        ),
+        ToolParam(
+            input_schema=Keyboard.model_json_schema(),
+            name="Keyboard",
+            description="""
+- Interact with docker container running image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
+- If user hasn't provided docker image id, check using `docker ps` and provide the id.
+- Emulate keyboard input to the screen
+- Uses xdootool to send keyboard input, keys like Return, BackSpace, Escape, Page_Up, etc. can be used.
+- Do not use it to interact with Bash tool.
 """,
         ),
     ]
@@ -358,7 +402,7 @@ System information:
                                 }
                             )
                             try:
-                                output_or_done, _ = get_tool_output(
+                                output_or_dones, _ = get_tool_output(
                                     tool_parsed,
                                     enc,
                                     limit - cost,
@@ -366,45 +410,46 @@ System information:
                                     max_tokens=8000,
                                 )
                             except Exception as e:
-                                output_or_done = (
-                                    f"GOT EXCEPTION while calling tool. Error: {e}"
-                                )
+                                output_or_dones = [
+                                    (f"GOT EXCEPTION while calling tool. Error: {e}")
+                                ]
                                 tb = traceback.format_exc()
-                                error_console.print(output_or_done + "\n" + tb)
-
-                            if isinstance(output_or_done, DoneFlag):
-                                system_console.print(
-                                    f"\n# Task marked done, with output {output_or_done.task_output}",
-                                )
-                                return output_or_done.task_output, cost
-
-                            output = output_or_done
-                            if isinstance(output, ImageData):
-                                tool_results.append(
-                                    ToolResultBlockParam(
-                                        type="tool_result",
-                                        tool_use_id=tc["id"],
-                                        content=[
-                                            {
-                                                "type": "image",
-                                                "source": {
-                                                    "type": "base64",
-                                                    "media_type": output.media_type,
-                                                    "data": output.data,
-                                                },
-                                            }
-                                        ],
+                                error_console.print(str(output_or_dones) + "\n" + tb)
+
+                            if any(isinstance(x, DoneFlag) for x in output_or_dones):
+                                return "", cost
+
+                            tool_results_content: list[
+                                TextBlockParam | ImageBlockParam
+                            ] = []
+                            for output in output_or_dones:
+                                assert not isinstance(output, DoneFlag)
+                                if isinstance(output, ImageData):
+                                    tool_results_content.append(
+                                        {
+                                            "type": "image",
+                                            "source": {
+                                                "type": "base64",
+                                                "media_type": output.media_type,
+                                                "data": output.data,
+                                            },
+                                        }
                                     )
-                                )
 
-                            else:
-                                tool_results.append(
-                                    ToolResultBlockParam(
-                                        type="tool_result",
-                                        tool_use_id=tc["id"],
-                                        content=output,
+                                else:
+                                    tool_results_content.append(
+                                        {
+                                            "type": "text",
+                                            "text": output,
+                                        },
                                     )
+                            tool_results.append(
+                                ToolResultBlockParam(
+                                    type="tool_result",
+                                    tool_use_id=tc["id"],
+                                    content=tool_results_content,
                                 )
+                            )
                         else:
                             _histories.append(
                                 {"role": "assistant", "content": full_response}

wcgw/client/computer_use.py

@@ -0,0 +1,416 @@
+"""Computer Use Tool for Anthropic API"""
+
+import base64
+import time
+import shlex
+import os
+from abc import ABCMeta, abstractmethod
+from dataclasses import dataclass, fields, replace
+from enum import StrEnum
+from typing import Any, Literal, TypedDict, Union, Optional
+from uuid import uuid4
+
+from anthropic.types.beta import BetaToolComputerUse20241022Param, BetaToolUnionParam
+from .sys_utils import command_run
+from ..types_ import (
+    Keyboard,
+    LeftClickDrag,
+    Mouse,
+    MouseMove,
+    ScreenShot,
+    GetScreenInfo,
+)
+
+
+# Constants
+OUTPUT_DIR = "/tmp/outputs"
+TYPING_DELAY_MS = 12
+TYPING_GROUP_SIZE = 50
+TRUNCATED_MESSAGE: str = "<response clipped><NOTE>To save on context only part of this file has been shown to you.</NOTE>"
+
+Action = Literal[
+    "key",
+    "type",
+    "mouse_move",
+    "left_click",
+    "left_click_drag",
+    "right_click",
+    "middle_click",
+    "double_click",
+    "screenshot",
+    "cursor_position",
+    "scroll_up",
+    "scroll_down",
+    "get_screen_info",
+]
+
+
+class Resolution(TypedDict):
+    width: int
+    height: int
+
+
+# Sizes above XGA/WXGA are not recommended
+MAX_SCALING_TARGETS: dict[str, Resolution] = {
+    "XGA": Resolution(width=1024, height=768),  # 4:3
+    "WXGA": Resolution(width=1280, height=800),  # 16:10
+    "FWXGA": Resolution(width=1366, height=768),  # ~16:9
+}
+
+
+class ScalingSource(StrEnum):
+    COMPUTER = "computer"
+    API = "api"
+
+
+class ComputerToolOptions(TypedDict):
+    display_height_px: int
+    display_width_px: int
+    display_number: int | None
+
+
+@dataclass(kw_only=True, frozen=True)
+class ToolResult:
+    """Represents the result of a tool execution."""
+
+    output: str | None = None
+    error: str | None = None
+    base64_image: str | None = None
+    system: str | None = None
+
+    def __bool__(self) -> bool:
+        return any(getattr(self, field.name) for field in fields(self))
+
+    def __add__(self, other: "ToolResult") -> "ToolResult":
+        def combine_fields(
+            field: str | None, other_field: str | None, concatenate: bool = True
+        ) -> str | None:
+            if field and other_field:
+                if concatenate:
+                    return field + other_field
+                raise ValueError("Cannot combine tool results")
+            return field or other_field
+
+        return ToolResult(
+            output=combine_fields(self.output, other.output),
+            error=combine_fields(self.error, other.error),
+            base64_image=combine_fields(self.base64_image, other.base64_image, False),
+            system=combine_fields(self.system, other.system),
+        )
+
+    def replace(self, **kwargs: Any) -> "ToolResult":
+        """Returns a new ToolResult with the given fields replaced."""
+        return replace(self, **kwargs)
+
+
+class CLIResult(ToolResult):
+    """A ToolResult that can be rendered as a CLI output."""
+
+    pass
+
+
+class ToolFailure(ToolResult):
+    """A ToolResult that represents a failure."""
+
+    pass
+
+
+class ToolError(Exception):
+    """Raised when a tool encounters an error."""
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+
+
+def chunks(s: str, chunk_size: int) -> list[str]:
+    return [s[i : i + chunk_size] for i in range(0, len(s), chunk_size)]
+
+
+class ComputerTool:
+    """
+    A tool that allows the agent to interact with the screen, keyboard, and mouse of the current computer.
+    The tool parameters are defined by Anthropic and are not editable.
+    """
+
+    name: Literal["computer"] = "computer"
+    api_type: Literal["computer_20241022"] = "computer_20241022"
+    width: Optional[int]
+    height: Optional[int]
+    display_num: Optional[int]
+    xdotool: Optional[str]
+    docker_image_id: Optional[str]
+
+    _screenshot_delay = 0.5
+    _scaling_enabled = True
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.xdotool = None
+        self.width = None
+        self.height = None
+        self.display_num = None
+        self._display_prefix = ""
+        self.docker_image_id = None
+
+    def get_screen_info(self) -> tuple[int, int, Optional[int]]:
+        result = self.shell(
+            "echo $WIDTH,$HEIGHT,$DISPLAY_NUM",
+            take_screenshot=False,
+        )
+        assert not result.error, result.error
+        assert result.output, "Could not get screen info"
+        width, height, display_num = map(
+            lambda x: None if not x else int(x), result.output.split(",")
+        )
+        if width is None:
+            width = 1080
+        if height is None:
+            height = 1920
+
+        self.width = width
+        self.height = height
+        if display_num is not None:
+            self.display_num = int(display_num)
+            self._display_prefix = f"DISPLAY=:{self.display_num} "
+        else:
+            self.display_num = None
+            self._display_prefix = ""
+        assert self._display_prefix is not None
+        self.xdotool = f"{self._display_prefix}xdotool"
+        return width, height, display_num
+
+    def __call__(
+        self,
+        *,
+        action: Action,
+        docker_image_id: Optional[str] = None,
+        text: str | None = None,
+        coordinate: tuple[int, int] | None = None,
+        **kwargs: Any,
+    ) -> ToolResult:
+        if action == "get_screen_info":
+            assert docker_image_id is not None
+            self.docker_image_id = docker_image_id
+            self.get_screen_info()
+            screenshot_res = self.screenshot()
+            return ToolResult(
+                output=f"width: {self.width}, height: {self.height}, display_num: {self.display_num}",
+                error=screenshot_res.error,
+                base64_image=screenshot_res.base64_image,
+            )
+
+        if self.width is None or self.height is None or self.docker_image_id is None:
+            raise ToolError("Please first get screen info using get_screen_info tool")
+
+        if action in ("mouse_move", "left_click_drag"):
+            if coordinate is None:
+                raise ToolError(f"coordinate is required for {action}")
+            if text is not None:
+                raise ToolError(f"text is not accepted for {action}")
+            if not isinstance(coordinate, (list, tuple)) or len(coordinate) != 2:
+                raise ToolError(f"{coordinate} must be a tuple of length 2")
+            if not all(isinstance(i, int) and i >= 0 for i in coordinate):
+                raise ToolError(f"{coordinate} must be a tuple of non-negative ints")
+
+            x, y = self.scale_coordinates(
+                ScalingSource.API, coordinate[0], coordinate[1]
+            )
+
+            if action == "mouse_move":
+                return self.shell(f"{self.xdotool} mousemove  {x} {y}")
+            elif action == "left_click_drag":
+                return self.shell(
+                    f"{self.xdotool} mousedown 1 mousemove  {x} {y} mouseup 1",
+                )
+
+        if action in ("key", "type"):
+            if text is None:
+                raise ToolError(f"text is required for {action}")
+            if coordinate is not None:
+                raise ToolError(f"coordinate is not accepted for {action}")
+            if not isinstance(text, str):
+                raise ToolError(output=f"{text} must be a string")
+
+            if action == "key":
+                return self.shell(f"{self.xdotool} key -- {text}")
+            elif action == "type":
+                results: list[ToolResult] = []
+                for chunk in chunks(text, TYPING_GROUP_SIZE):
+                    cmd = f"{self.xdotool} type --delay {TYPING_DELAY_MS} -- {shlex.quote(chunk)}"
+                    results.append(self.shell(cmd, take_screenshot=False))
+                screenshot_base64 = self.screenshot().base64_image
+                return ToolResult(
+                    output="".join(result.output or "" for result in results),
+                    error="".join(result.error or "" for result in results),
+                    base64_image=screenshot_base64,
+                )
+
+        if action in (
+            "left_click",
+            "right_click",
+            "double_click",
+            "middle_click",
+            "screenshot",
+            "cursor_position",
+            "scroll_up",
+            "scroll_down",
+        ):
+            if text is not None:
+                raise ToolError(f"text is not accepted for {action}")
+            if coordinate is not None:
+                raise ToolError(f"coordinate is not accepted for {action}")
+
+            if action == "screenshot":
+                return self.screenshot()
+            elif action == "cursor_position":
+                result = self.shell(
+                    f"{self.xdotool} getmouselocation --shell",
+                    take_screenshot=False,
+                )
+                output = result.output or ""
+                x, y = self.scale_coordinates(
+                    ScalingSource.COMPUTER,
+                    int(output.split("X=")[1].split("\n")[0]),
+                    int(output.split("Y=")[1].split("\n")[0]),
+                )
+                return result.replace(output=f"X={x},Y={y}")
+            else:
+                if action in ("scroll_up", "scroll_down"):
+                    button = "4" if action == "scroll_up" else "5"
+                    return self.shell(
+                        f"{self.xdotool} click --repeat 1 {button}",
+                    )
+                else:
+                    click_arg = {
+                        "left_click": "1",
+                        "right_click": "3",
+                        "middle_click": "2",
+                        "double_click": "--repeat 2 --delay 500 1",
+                    }[action]
+                    return self.shell(f"{self.xdotool} click {click_arg}")
+
+        raise ToolError(f"Invalid action: {action}")
+
+    def screenshot(self) -> ToolResult:
+        """Take a screenshot of the current screen and return the base64 encoded image."""
+        if self.width is None or self.height is None or self.docker_image_id is None:
+            self.get_screen_info()
+        assert self.width and self.height
+        # output_dir = Path(OUTPUT_DIR)
+        # output_dir.mkdir(parents=True, exist_ok=True)
+        mkdir_res = self.shell(
+            command=f"mkdir -p {OUTPUT_DIR}",
+            take_screenshot=False,
+        )
+        path = f"{OUTPUT_DIR}/screenshot_{uuid4().hex}.png"
+
+        screenshot_cmd = f"{self._display_prefix}scrot -f {path} -p"
+
+        self.shell(screenshot_cmd, take_screenshot=False)
+
+        if self._scaling_enabled:
+            x, y = self.scale_coordinates(
+                ScalingSource.COMPUTER, self.width, self.height
+            )
+            self.shell(
+                f"convert {path} -resize {x}x{y}! {path}",
+                take_screenshot=False,
+            )
+
+        # Copy file from docker to tmp
+        _, stdout, stderr = command_run(
+            f"docker cp {self.docker_image_id}:{path} {path}",
+            truncate_after=None,
+        )
+
+        if os.path.exists(path):
+            with open(path, "rb") as f:
+                base64_image = base64.b64encode(f.read()).decode("utf-8")
+
+            return ToolResult(output="", error=stderr, base64_image=base64_image)
+
+        raise ToolError(f"Failed to take screenshot: {stderr}")
+
+    def shell(self, command: str, take_screenshot: bool = True) -> ToolResult:
+        """Run a shell command and return the output, error, and optionally a screenshot."""
+        _, stdout, stderr = command_run(
+            f"docker exec {self.docker_image_id} sh -c '{command}'"
+        )
+        base64_image = None
+
+        if take_screenshot:
+            # delay to let things settle before taking a screenshot
+            time.sleep(self._screenshot_delay)
+            base64_image = self.screenshot().base64_image
+
+        return ToolResult(output=stdout, error=stderr, base64_image=base64_image)
+
+    def scale_coordinates(
+        self, source: ScalingSource, x: int, y: int
+    ) -> tuple[int, int]:
+        """Scale coordinates to a target maximum resolution."""
+
+        if self.width is None or self.height is None:
+            raise ToolError("Please first get screen info using get_screen_info tool")
+
+        if not self._scaling_enabled:
+            return x, y
+        ratio = self.width / self.height
+        target_dimension = None
+        for dimension in MAX_SCALING_TARGETS.values():
+            # allow some error in the aspect ratio - not ratios are exactly 16:9
+            if abs(dimension["width"] / dimension["height"] - ratio) < 0.02:
+                if dimension["width"] < self.width:
+                    target_dimension = dimension
+                break
+        if target_dimension is None:
+            return x, y
+        # should be less than 1
+        x_scaling_factor = target_dimension["width"] / self.width
+        y_scaling_factor = target_dimension["height"] / self.height
+        if source == ScalingSource.API:
+            if x > self.width or y > self.height:
+                raise ToolError(f"Coordinates {x}, {y} are out of bounds")
+            # scale up
+            return round(x / x_scaling_factor), round(y / y_scaling_factor)
+        # scale down
+        return round(x * x_scaling_factor), round(y * y_scaling_factor)
+
+
+Computer = ComputerTool()
+
+
+def run_computer_tool(
+    action: Union[Keyboard, Mouse, ScreenShot, GetScreenInfo],
+) -> tuple[str, str]:
+    if isinstance(action, GetScreenInfo):
+        result = Computer(
+            action="get_screen_info", docker_image_id=action.docker_image_id
+        )
+    elif isinstance(action, ScreenShot):
+        result = Computer(
+            action="screenshot",
+        )
+    elif isinstance(action, Keyboard):
+        result = Computer(
+            action=action.action,
+            text=action.text,
+        )
+    elif isinstance(action, Mouse):
+        if isinstance(action.action, MouseMove):
+            result = Computer(
+                action="mouse_move",
+                coordinate=(action.action.x, action.action.y),
+            )
+        elif isinstance(action.action, LeftClickDrag):
+            result = Computer(
+                action="left_click_drag",
+                coordinate=(action.action.x, action.action.y),
+            )
+        else:
+            result = Computer(action=action.action.button_type)
+
+    output = f"stdout: {result.output or ''}, stderr: {result.error or ''}"
+    image = result.base64_image or ""
+    return output, image

wcgw/client/mcp_server/Readme.md

@@ -0,0 +1,73 @@
+# Claude desktop support
+
+## Setup
+
+Install xdtool
+
+```sh
+brew install xdotool
+
+# On macos:
+defaults write org.x.X11 enable_test_extensions -boolean true
+```
+
+Update `claude_desktop_config.json` (~/Library/Application Support/Claude/claude_desktop_config.json)
+
+```json
+{
+  "mcpServers": {
+    "wcgw": {
+      "command": "uvx",
+      "args": ["--from", "wcgw@latest", "wcgw_mcp"]
+    }
+  }
+}
+```
+
+Then restart claude app.
+
+### Computer use support using desktop on docker
+
+Controlling the system isn't possible yet, but you can connect to a docker container which runs a linux os with desktop.
+
+First run a sample docker image with desktop and optionally VNC connection:
+
+```sh
+docker run \
+    --entrypoint "" \
+    -p 6080:6080 \
+    -e WIDTH=1024 \
+    -e HEIGHT=768 \
+    -d \
+    ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest \
+    bash -c "\
+        ./start_all.sh && \
+        ./novnc_startup.sh && \
+        python http_server.py > /tmp/server_logs.txt 2>&1 & \
+        tail -f /dev/null"
+```
+
+Connect to `http://localhost:6080/vnc.html` for desktop view (VNC) of the system running in the docker.
+
+## Usage
+
+Wait for a few seconds. You should be able to see this icon if everything goes right.
+
+![mcp icon](https://github.com/rusiaaman/wcgw/blob/main/static/rocket-icon.png?raw=true)
+over here
+
+![mcp icon](https://github.com/rusiaaman/wcgw/blob/main/static/claude-ss.jpg?raw=true)
+
+Then ask claude to execute shell commands, read files, edit files, run your code, etc.
+
+If you've run the docker for LLM to access, you can ask it to control the "docker os". If you don't provide the docker container id to it, it'll try to search for available docker using `docker ps` command.
+
+## Example
+
+### Computer use example
+
+![computer-use](https://github.com/rusiaaman/wcgw/blob/main/static/computer-use.jpg?raw=true)
+
+### Shell example
+
+![example](https://github.com/rusiaaman/wcgw/blob/main/static/example.jpg?raw=true)