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

wcgw/client/tools.py

@@ -11,6 +11,7 @@ import threading
 import importlib.metadata
 import time
 import traceback
+from tempfile import TemporaryDirectory
 from typing import (
     Callable,
     Literal,
@@ -24,6 +25,7 @@ from typing import (
 import uuid
 from pydantic import BaseModel, TypeAdapter
 import typer
+from .computer_use import run_computer_tool
 from websockets.sync.client import connect as syncconnect
 
 import os
@@ -57,11 +59,14 @@ from ..types_ import (
     ReadFile,
     ReadImage,
     ResetShell,
-    Writefile,
+    Mouse,
+    Keyboard,
+    ScreenShot,
+    GetScreenInfo,
 )
 
 from .common import CostData, Models, discard_input
-
+from .sys_utils import command_run
 from .openai_utils import get_input_cost, get_output_cost
 
 console = rich.console.Console(style="magenta", highlight=False, markup=False)
@@ -96,7 +101,8 @@ def ask_confirmation(prompt: Confirmation) -> str:
     return "Yes" if response.lower() == "y" else "No"
 
 
-PROMPT = "#@@"
+PROMPT_CONST = "#@wcgw@#"
+PROMPT = PROMPT_CONST
 
 
 def start_shell() -> pexpect.spawn:  # type: ignore
@@ -125,7 +131,7 @@ def _is_int(mystr: str) -> bool:
 
 
 def _get_exit_code() -> int:
-    if PROMPT != "#@@":
+    if PROMPT != PROMPT_CONST:
         return 0
     # First reset the prompt in case venv was sourced or other reasons.
     SHELL.sendline(f"export PS1={PROMPT}")
@@ -153,6 +159,7 @@ def _get_exit_code() -> int:
 
 BASH_CLF_OUTPUT = Literal["repl", "pending"]
 BASH_STATE: BASH_CLF_OUTPUT = "repl"
+IS_IN_DOCKER: Optional[str] = ""
 CWD = os.getcwd()
 
 
@@ -167,10 +174,11 @@ Current working directory: {CWD}
 
 
 def reset_shell() -> str:
-    global SHELL, BASH_STATE, CWD
+    global SHELL, BASH_STATE, CWD, IS_IN_DOCKER
     SHELL.close(True)
     SHELL = start_shell()
     BASH_STATE = "repl"
+    IS_IN_DOCKER = ""
     CWD = os.getcwd()
     return "Reset successful" + get_status()
 
@@ -234,6 +242,7 @@ def execute_bash(
     enc: tiktoken.Encoding,
     bash_arg: BashCommand | BashInteraction,
     max_tokens: Optional[int],
+    timeout_s: Optional[float],
 ) -> tuple[str, float]:
     global SHELL, BASH_STATE, CWD
     try:
@@ -333,7 +342,7 @@ def execute_bash(
         SHELL.expect(PROMPT)
         return "---\n\nFailure: user interrupted the execution", 0.0
 
-    wait = 5
+    wait = timeout_s or 5
     index = SHELL.expect([PROMPT, pexpect.TIMEOUT], timeout=wait)
     if index == 1:
         BASH_STATE = "pending"
@@ -438,35 +447,70 @@ def read_image_from_shell(file_path: str) -> ImageData:
     if not os.path.isabs(file_path):
         file_path = os.path.join(CWD, file_path)
 
-    if not os.path.exists(file_path):
-        raise ValueError(f"File {file_path} does not exist")
-
-    with open(file_path, "rb") as image_file:
-        image_bytes = image_file.read()
-        image_b64 = base64.b64encode(image_bytes).decode("utf-8")
-        image_type = mimetypes.guess_type(file_path)[0]
-        return ImageData(media_type=image_type, data=image_b64)  # type: ignore
-
+    if not IS_IN_DOCKER:
+        if not os.path.exists(file_path):
+            raise ValueError(f"File {file_path} does not exist")
 
-def write_file(writefile: Writefile | CreateFileNew, error_on_exist: bool) -> str:
+        with open(file_path, "rb") as image_file:
+            image_bytes = image_file.read()
+            image_b64 = base64.b64encode(image_bytes).decode("utf-8")
+            image_type = mimetypes.guess_type(file_path)[0]
+            return ImageData(media_type=image_type, data=image_b64)  # type: ignore
+    else:
+        with TemporaryDirectory() as tmpdir:
+            rcode = os.system(f"docker cp {IS_IN_DOCKER}:{file_path} {tmpdir}")
+            if rcode != 0:
+                raise Exception(f"Error: Read failed with code {rcode}")
+            path_ = os.path.join(tmpdir, os.path.basename(file_path))
+            with open(path_, "rb") as f:
+                image_bytes = f.read()
+            image_b64 = base64.b64encode(image_bytes).decode("utf-8")
+            image_type = mimetypes.guess_type(file_path)[0]
+            return ImageData(media_type=image_type, data=image_b64)  # type: ignore
+
+
+def write_file(writefile: CreateFileNew, error_on_exist: bool) -> str:
     if not os.path.isabs(writefile.file_path):
         return f"Failure: file_path should be absolute path, current working directory is {CWD}"
     else:
         path_ = writefile.file_path
 
-    if error_on_exist and os.path.exists(path_):
-        file_data = Path(path_).read_text()
-        if file_data:
-            return f"Error: can't write to existing file {path_}, use other functions to edit the file"
+    if not IS_IN_DOCKER:
+        if error_on_exist and os.path.exists(path_):
+            file_data = Path(path_).read_text()
+            if file_data:
+                return f"Error: can't write to existing file {path_}, use other functions to edit the file"
 
-    path = Path(path_)
-    path.parent.mkdir(parents=True, exist_ok=True)
+        path = Path(path_)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            with path.open("w") as f:
+                f.write(writefile.file_content)
+        except OSError as e:
+            return f"Error: {e}"
+    else:
+        if error_on_exist:
+            # Check if it exists using os.system
+            cmd = f"test -f {path_}"
+            status = os.system(f'docker exec {IS_IN_DOCKER} bash -c "{cmd}"')
+            if status == 0:
+                return f"Error: can't write to existing file {path_}, use other functions to edit the file"
+
+        with TemporaryDirectory() as tmpdir:
+            tmppath = os.path.join(tmpdir, os.path.basename(path_))
+            with open(tmppath, "w") as f:
+                f.write(writefile.file_content)
+            os.chmod(tmppath, 0o777)
+            parent_dir = os.path.dirname(path_)
+            rcode = os.system(f"docker exec {IS_IN_DOCKER} mkdir -p {parent_dir}")
+            if rcode != 0:
+                return f"Error: Write failed with code while creating dirs {rcode}"
+
+            rcode = os.system(f"docker cp {tmppath} {IS_IN_DOCKER}:{path_}")
+            if rcode != 0:
+                return f"Error: Write failed with code {rcode}"
 
-    try:
-        with path.open("w") as f:
-            f.write(writefile.file_content)
-    except OSError as e:
-        return f"Error: {e}"
     console.print(f"File written to {path_}")
     return "Success"
 
@@ -545,11 +589,21 @@ def do_diff_edit(fedit: FileEdit) -> str:
     else:
         path_ = fedit.file_path
 
-    if not os.path.exists(path_):
-        raise Exception(f"Error: file {path_} does not exist")
+    if not IS_IN_DOCKER:
+        if not os.path.exists(path_):
+            raise Exception(f"Error: file {path_} does not exist")
 
-    with open(path_) as f:
-        apply_diff_to = f.read()
+        with open(path_) as f:
+            apply_diff_to = f.read()
+    else:
+        # Copy from docker
+        with TemporaryDirectory() as tmpdir:
+            rcode = os.system(f"docker cp {IS_IN_DOCKER}:{path_} {tmpdir}")
+            if rcode != 0:
+                raise Exception(f"Error: Read failed with code {rcode}")
+            path_tmp = os.path.join(tmpdir, os.path.basename(path_))
+            with open(path_tmp, "r") as f:
+                apply_diff_to = f.read()
 
     fedit.file_edit_using_search_replace_blocks = (
         fedit.file_edit_using_search_replace_blocks.strip()
@@ -597,40 +651,20 @@ def do_diff_edit(fedit: FileEdit) -> str:
             "Error: no valid search-replace blocks found, please check your syntax for FileEdit"
         )
 
-    with open(path_, "w") as f:
-        f.write(apply_diff_to)
-
-    return "Success"
-
-
-def file_edit(fedit: FileEditFindReplace) -> str:
-    if not os.path.isabs(fedit.file_path):
-        raise Exception(
-            f"Failure: file_path should be absolute path, current working directory is {CWD}"
-        )
+    if not IS_IN_DOCKER:
+        with open(path_, "w") as f:
+            f.write(apply_diff_to)
     else:
-        path_ = fedit.file_path
+        with TemporaryDirectory() as tmpdir:
+            path_tmp = os.path.join(tmpdir, os.path.basename(path_))
+            with open(path_tmp, "w") as f:
+                f.write(apply_diff_to)
+            os.chmod(path_tmp, 0o777)
+            # Copy to docker using docker cp
+            rcode = os.system(f"docker cp {path_tmp} {IS_IN_DOCKER}:{path_}")
+            if rcode != 0:
+                raise Exception(f"Error: Write failed with code {rcode}")
 
-    if not os.path.exists(path_):
-        raise Exception(f"Error: file {path_} does not exist")
-
-    if not fedit.find_lines:
-        raise Exception("Error: `find_lines` cannot be empty")
-
-    out_string = "\n".join("> " + line for line in fedit.find_lines.split("\n"))
-    in_string = "\n".join("< " + line for line in fedit.replace_with_lines.split("\n"))
-    console.log(f"Editing file: {path_}\n---\n{out_string}\n---\n{in_string}\n---")
-    try:
-        with open(path_) as f:
-            content = f.read()
-
-        content = edit_content(content, fedit.find_lines, fedit.replace_with_lines)
-
-        with open(path_, "w") as f:
-            f.write(content)
-    except OSError as e:
-        raise Exception(f"Error: {e}")
-    console.print(f"File written to {path_}")
     return "Success"
 
 
@@ -661,7 +695,6 @@ TOOLS = (
     | BashCommand
     | BashInteraction
     | ResetShell
-    | Writefile
     | CreateFileNew
     | FileEditFindReplace
     | FileEdit
@@ -670,6 +703,10 @@ TOOLS = (
     | ReadImage
     | ReadFile
     | Initialize
+    | Mouse
+    | Keyboard
+    | ScreenShot
+    | GetScreenInfo
 )
 
 
@@ -687,8 +724,6 @@ def which_tool_name(name: str) -> Type[TOOLS]:
         return BashInteraction
     elif name == "ResetShell":
         return ResetShell
-    elif name == "Writefile":
-        return Writefile
     elif name == "CreateFileNew":
         return CreateFileNew
     elif name == "FileEditFindReplace":
@@ -705,6 +740,14 @@ def which_tool_name(name: str) -> Type[TOOLS]:
         return ReadFile
     elif name == "Initialize":
         return Initialize
+    elif name == "Mouse":
+        return Mouse
+    elif name == "Keyboard":
+        return Keyboard
+    elif name == "ScreenShot":
+        return ScreenShot
+    elif name == "GetScreenInfo":
+        return GetScreenInfo
     else:
         raise ValueError(f"Unknown tool name: {name}")
 
@@ -715,7 +758,6 @@ def get_tool_output(
     | BashCommand
     | BashInteraction
     | ResetShell
-    | Writefile
     | CreateFileNew
     | FileEditFindReplace
     | FileEdit
@@ -723,19 +765,23 @@ def get_tool_output(
     | DoneFlag
     | ReadImage
     | Initialize
-    | ReadFile,
+    | ReadFile
+    | Mouse
+    | Keyboard
+    | ScreenShot
+    | GetScreenInfo,
     enc: tiktoken.Encoding,
     limit: float,
     loop_call: Callable[[str, float], tuple[str, float]],
     max_tokens: Optional[int],
-) -> tuple[str | ImageData | DoneFlag, float]:
+) -> tuple[list[str | ImageData | DoneFlag], float]:
+    global IS_IN_DOCKER
     if isinstance(args, dict):
         adapter = TypeAdapter[
             Confirmation
             | BashCommand
             | BashInteraction
             | ResetShell
-            | Writefile
             | CreateFileNew
             | FileEditFindReplace
             | FileEdit
@@ -744,12 +790,15 @@ def get_tool_output(
             | ReadImage
             | ReadFile
             | Initialize
+            | Mouse
+            | Keyboard
+            | ScreenShot
+            | GetScreenInfo,
         ](
             Confirmation
             | BashCommand
             | BashInteraction
             | ResetShell
-            | Writefile
             | CreateFileNew
             | FileEditFindReplace
             | FileEdit
@@ -758,6 +807,10 @@ def get_tool_output(
             | ReadImage
             | ReadFile
             | Initialize
+            | Mouse
+            | Keyboard
+            | ScreenShot
+            | GetScreenInfo
         )
         arg = adapter.validate_python(args)
     else:
@@ -768,16 +821,10 @@ def get_tool_output(
         output = ask_confirmation(arg), 0.0
     elif isinstance(arg, (BashCommand | BashInteraction)):
         console.print("Calling execute bash tool")
-        output = execute_bash(enc, arg, max_tokens)
-    elif isinstance(arg, Writefile):
-        console.print("Calling write file tool")
-        output = write_file(arg, False), 0
+        output = execute_bash(enc, arg, max_tokens, None)
     elif isinstance(arg, CreateFileNew):
         console.print("Calling write file tool")
         output = write_file(arg, True), 0
-    elif isinstance(arg, FileEditFindReplace):
-        console.print("Calling file edit tool")
-        output = file_edit(arg), 0.0
     elif isinstance(arg, FileEdit):
         console.print("Calling full file edit tool")
         output = do_diff_edit(arg), 0.0
@@ -799,11 +846,50 @@ def get_tool_output(
     elif isinstance(arg, Initialize):
         console.print("Calling initial info tool")
         output = initial_info(), 0.0
+    elif isinstance(arg, (Mouse, Keyboard, ScreenShot, GetScreenInfo)):
+        console.print(f"Calling {type(arg).__name__} tool")
+        outputs_cost = run_computer_tool(arg), 0.0
+        console.print(outputs_cost[0][0])
+        outputs: list[ImageData | str | DoneFlag] = [outputs_cost[0][0]]
+        imgBs64 = outputs_cost[0][1]
+        if imgBs64:
+            console.print("Captured screenshot")
+            outputs.append(ImageData(media_type="image/png", data=imgBs64))
+            if not IS_IN_DOCKER:
+                try:
+                    # At this point we should go into the docker env
+                    res, _ = execute_bash(
+                        enc,
+                        BashCommand(
+                            command=f"docker exec -it {arg.docker_image_id} sh"
+                        ),
+                        None,
+                        0.2,
+                    )
+                    # At this point we should go into the docker env
+                    res, _ = execute_bash(
+                        enc,
+                        BashInteraction(
+                            send_text=f"export PS1={PROMPT}", type="BashInteraction"
+                        ),
+                        None,
+                        0.2,
+                    )
+                    # Do chown of home dir
+                except Exception as e:
+                    reset_shell()
+                    raise Exception(
+                        f"Some error happened while going inside docker. I've reset the shell. Please start again. Error {e}"
+                    )
+                IS_IN_DOCKER = arg.docker_image_id
+        return outputs, outputs_cost[1]
     else:
         raise ValueError(f"Unknown tool: {arg}")
-
-    console.print(str(output[0]))
-    return output
+    if isinstance(output[0], str):
+        console.print(str(output[0]))
+    else:
+        console.print(f"Received {type(output[0])} from tool")
+    return [output[0]], output[1]
 
 
 History = list[ChatCompletionMessageParam]
@@ -818,7 +904,6 @@ class Mdata(BaseModel):
     data: (
         BashCommand
         | BashInteraction
-        | Writefile
         | CreateFileNew
         | ResetShell
         | FileEditFindReplace
@@ -853,9 +938,10 @@ def register_client(server_url: str, client_uuid: str = "") -> None:
                 if isinstance(mdata.data, str):
                     raise Exception(mdata)
                 try:
-                    output, cost = get_tool_output(
+                    outputs, cost = get_tool_output(
                         mdata.data, default_enc, 0.0, lambda x, y: ("", 0), 8000
                     )
+                    output = outputs[0]
                     curr_cost += cost
                     print(f"{curr_cost=}")
                 except Exception as e:
@@ -893,12 +979,22 @@ def read_file(readfile: ReadFile, max_tokens: Optional[int]) -> str:
     if not os.path.isabs(readfile.file_path):
         return f"Failure: file_path should be absolute path, current working directory is {CWD}"
 
-    path = Path(readfile.file_path)
-    if not path.exists():
-        return f"Error: file {readfile.file_path} does not exist"
+    if not IS_IN_DOCKER:
+        path = Path(readfile.file_path)
+        if not path.exists():
+            return f"Error: file {readfile.file_path} does not exist"
 
-    with path.open("r") as f:
-        content = f.read()
+        with path.open("r") as f:
+            content = f.read()
+
+    else:
+        return_code, content, stderr = command_run(
+            f"cat {readfile.file_path}",
+        )
+        if return_code != 0:
+            raise Exception(
+                f"Error: cat {readfile.file_path} failed with code {return_code}\nstdout: {content}\nstderr: {stderr}"
+            )
 
     if max_tokens is not None:
         tokens = default_enc.encode(content)

wcgw/relay/serve.py

@@ -23,7 +23,6 @@ from ..types_ import (
     Initialize,
     ReadFile,
     ResetShell,
-    Writefile,
     Specials,
 )
 
@@ -32,7 +31,6 @@ class Mdata(BaseModel):
     data: (
         BashCommand
         | BashInteraction
-        | Writefile
         | CreateFileNew
         | ResetShell
         | FileEditFindReplace
@@ -67,8 +65,7 @@ async def register_websocket(websocket: WebSocket, uuid: UUID) -> None:
     # receive client version
     client_version = await websocket.receive_text()
     sem_version_client = semantic_version.Version.coerce(client_version)
-    sem_version_server = semantic_version.Version.coerce(
-        CLIENT_VERSION_MINIMUM)
+    sem_version_server = semantic_version.Version.coerce(CLIENT_VERSION_MINIMUM)
     if sem_version_client < sem_version_server:
         await websocket.send_text(
             Mdata(
@@ -93,8 +90,7 @@ async def register_websocket(websocket: WebSocket, uuid: UUID) -> None:
         while True:
             received_data = await websocket.receive_text()
             if uuid not in gpts:
-                raise fastapi.HTTPException(
-                    status_code=400, detail="No call made")
+                raise fastapi.HTTPException(status_code=400, detail="No call made")
             gpts[uuid](received_data)
     except WebSocketDisconnect:
         # Remove the client if the WebSocket is disconnected
@@ -281,9 +277,7 @@ async def read_file_endpoint(read_file_data: ReadFileWithUUID) -> str:
 
     gpts[user_id] = put_results
 
-    await clients[user_id](
-        Mdata(data=read_file_data, user_id=user_id)
-    )
+    await clients[user_id](Mdata(data=read_file_data, user_id=user_id))
 
     start_time = time.time()
     while time.time() - start_time < 30:
@@ -293,6 +287,7 @@ async def read_file_endpoint(read_file_data: ReadFileWithUUID) -> str:
 
     raise fastapi.HTTPException(status_code=500, detail="Timeout error")
 
+
 class InitializeWithUUID(Initialize):
     user_id: UUID
 
@@ -311,9 +306,7 @@ async def initialize(initialize_data: InitializeWithUUID) -> str:
 
     gpts[user_id] = put_results
 
-    await clients[user_id](
-        Mdata(data=initialize_data, user_id=user_id)
-    )
+    await clients[user_id](Mdata(data=initialize_data, user_id=user_id))
 
     start_time = time.time()
     while time.time() - start_time < 30:

wcgw/types_.py

@@ -24,11 +24,6 @@ class ReadImage(BaseModel):
     type: Literal["ReadImage"]
 
 
-class Writefile(BaseModel):
-    file_path: str
-    file_content: str
-
-
 class CreateFileNew(BaseModel):
     file_path: str
     file_content: str
@@ -55,4 +50,45 @@ class FileEdit(BaseModel):
 
 
 class Initialize(BaseModel):
-    type: Literal["Initialize"]
+    type: Literal["Initialize"]
+
+
+class GetScreenInfo(BaseModel):
+    type: Literal["GetScreenInfo"]
+    docker_image_id: str
+
+
+class ScreenShot(BaseModel):
+    type: Literal["ScreenShot"]
+    docker_image_id: str
+
+
+class MouseMove(BaseModel):
+    x: int
+    y: int
+    type: Literal["MouseMove"]
+
+
+class LeftClickDrag(BaseModel):
+    x: int
+    y: int
+
+
+class MouseButton(BaseModel):
+    button_type: Literal[
+        "left_click",
+        "right_click",
+        "middle_click",
+        "double_click",
+        "scroll_up",
+        "scroll_down",
+    ]
+
+
+class Mouse(BaseModel):
+    action: MouseButton | LeftClickDrag | MouseMove
+
+
+class Keyboard(BaseModel):
+    action: Literal["key", "type"]
+    text: str