wcgw 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

wcgw/client/tools.py

@@ -1,7 +1,10 @@
 import asyncio
 import base64
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from io import BytesIO
 import json
 import mimetypes
+from pathlib import Path
 import re
 import sys
 import threading
@@ -13,6 +16,7 @@ from typing import (
     NewType,
     Optional,
     ParamSpec,
+    Type,
     TypeVar,
     TypedDict,
 )
@@ -41,7 +45,14 @@ from openai.types.chat import (
     ParsedChatCompletionMessage,
 )
 from nltk.metrics.distance import edit_distance
-from ..types_ import FileEditFindReplace, ResetShell, Writefile
+
+from ..types_ import (
+    CreateFileNew,
+    FileEditFindReplace,
+    FullFileEdit,
+    ResetShell,
+    Writefile,
+)
 
 from ..types_ import BashCommand
 
@@ -70,7 +81,10 @@ def render_terminal_output(text: str) -> str:
             break
     else:
         i = len(dsp)
-    return "\n".join(screen.display[: len(dsp) - i])
+    lines = screen.display[: len(dsp) - i]
+    # Strip trailing space
+    lines = [line.rstrip() for line in lines]
+    return "\n".join(lines)
 
 
 class Confirmation(BaseModel):
@@ -368,8 +382,16 @@ def serve_image_in_bg(file_path: str, client_uuid: str, name: str) -> None:
             serve_image_in_bg(file_path, client_uuid, name)
 
 
+MEDIA_TYPES = Literal["image/jpeg", "image/png", "image/gif", "image/webp"]
+
+
 class ImageData(BaseModel):
-    dataurl: str
+    media_type: MEDIA_TYPES
+    data: str
+
+    @property
+    def dataurl(self) -> str:
+        return f"data:{self.media_type};base64," + self.data
 
 
 Param = ParamSpec("Param")
@@ -399,16 +421,25 @@ def read_image_from_shell(file_path: str) -> ImageData:
         image_bytes = image_file.read()
         image_b64 = base64.b64encode(image_bytes).decode("utf-8")
         image_type = mimetypes.guess_type(file_path)[0]
-        return ImageData(dataurl=f"data:{image_type};base64,{image_b64}")
+        return ImageData(media_type=image_type, data=image_b64)
 
 
-def write_file(writefile: Writefile) -> str:
+def write_file(writefile: Writefile | CreateFileNew, error_on_exist: bool) -> str:
     if not os.path.isabs(writefile.file_path):
-        path_ = os.path.join(CWD, writefile.file_path)
+        return "Failure: file_path should be absolute path"
     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"
+
+    path = Path(path_)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
     try:
-        with open(path_, "w") as f:
+        with path.open("w") as f:
             f.write(writefile.file_content)
     except OSError as e:
         return f"Error: {e}"
@@ -416,50 +447,128 @@ def write_file(writefile: Writefile) -> str:
     return "Success"
 
 
-def find_least_edit_distance_substring(content: str, find_str: str) -> str:
-    content_lines = content.split("\n")
+def find_least_edit_distance_substring(
+    content: str, find_str: str
+) -> tuple[str, float]:
+    orig_content_lines = content.split("\n")
+    content_lines = [
+        line.strip() for line in orig_content_lines
+    ]  # Remove trailing and leading space for calculating edit distance
     find_lines = find_str.split("\n")
+    find_lines = [
+        line.strip() for line in find_lines
+    ]  # Remove trailing and leading space for calculating edit distance
     # Slide window and find one with sum of edit distance least
     min_edit_distance = float("inf")
     min_edit_distance_lines = []
-    for i in range(len(content_lines) - len(find_lines) + 1):
+    for i in range(max(1, len(content_lines) - len(find_lines) + 1)):
         edit_distance_sum = 0
         for j in range(len(find_lines)):
-            edit_distance_sum += edit_distance(content_lines[i + j], find_lines[j])
+            if (i + j) < len(content_lines):
+                edit_distance_sum += edit_distance(content_lines[i + j], find_lines[j])
+            else:
+                edit_distance_sum += len(find_lines[j])
         if edit_distance_sum < min_edit_distance:
             min_edit_distance = edit_distance_sum
-            min_edit_distance_lines = content_lines[i : i + len(find_lines)]
-    return "\n".join(min_edit_distance_lines)
+            min_edit_distance_lines = orig_content_lines[i : i + len(find_lines)]
+    return "\n".join(min_edit_distance_lines), min_edit_distance
+
+
+def edit_content(content: str, find_lines: str, replace_with_lines: str) -> str:
+    count = content.count(find_lines)
+    if count == 0:
+        closest_match, min_edit_distance = find_least_edit_distance_substring(
+            content, find_lines
+        )
+        print(
+            f"Exact match not found, found with whitespace removed edit distance: {min_edit_distance}"
+        )
+        if min_edit_distance / len(find_lines) < 1 / 100:
+            print("Editing file with closest match")
+            return edit_content(content, closest_match, replace_with_lines)
+        raise Exception(
+            f"Error: no match found for the provided `find_lines` in the file. Closest match:\n---\n{closest_match}\n---\nFile not edited"
+        )
 
+    content = content.replace(find_lines, replace_with_lines, 1)
+    return content
 
-def file_edit(file_edit: FileEditFindReplace) -> str:
-    if not os.path.isabs(file_edit.file_path):
-        path_ = os.path.join(CWD, file_edit.file_path)
+
+def do_diff_edit(fedit: FullFileEdit) -> str:
+    console.log(f"Editing file: {fedit.file_path}")
+
+    if not os.path.isabs(fedit.file_path):
+        raise Exception("Failure: file_path should be absolute path")
     else:
-        path_ = file_edit.file_path
+        path_ = fedit.file_path
+
+    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()
+
+    lines = fedit.file_edit_using_searh_replace_blocks.split("\n")
+    n_lines = len(lines)
+    i = 0
+    while i < n_lines:
+        if re.match(r"^<<<<<<+\s*SEARCH\s*$", lines[i]):
+            search_block = []
+            i += 1
+            while i < n_lines and not re.match(r"^======*\s*$", lines[i]):
+                search_block.append(lines[i])
+                i += 1
+            i += 1
+            replace_block = []
+            while i < n_lines and not re.match(r"^>>>>>>+\s*REPLACE\s*$", lines[i]):
+                replace_block.append(lines[i])
+                i += 1
+            i += 1
+
+            for line in search_block:
+                console.log("> " + line)
+            console.log("---")
+            for line in replace_block:
+                console.log("< " + line)
+
+            search_block_ = "\n".join(search_block)
+            replace_block_ = "\n".join(replace_block)
+
+            apply_diff_to = edit_content(apply_diff_to, search_block_, replace_block_)
+        else:
+            i += 1
 
-    out_string = "\n".join("> " + line for line in file_edit.find_lines.split("\n"))
-    in_string = "\n".join(
-        "< " + line for line in file_edit.replace_with_lines.split("\n")
-    )
+    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("Failure: file_path should be absolute path")
+    else:
+        path_ = fedit.file_path
+
+    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()
-        # First find counts
-        count = content.count(file_edit.find_lines)
 
-        if count == 0:
-            closest_match = find_least_edit_distance_substring(
-                content, file_edit.find_lines
-            )
-            return f"Error: no match found for the provided `find_lines` in the file. Closest match:\n---\n{closest_match}\n---\nFile not edited"
+        content = edit_content(content, fedit.find_lines, fedit.replace_with_lines)
 
-        content = content.replace(file_edit.find_lines, file_edit.replace_with_lines)
         with open(path_, "w") as f:
             f.write(content)
     except OSError as e:
-        return f"Error: {e}"
+        raise Exception(f"Error: {e}")
     console.print(f"File written to {path_}")
     return "Success"
 
@@ -486,31 +595,53 @@ def take_help_of_ai_assistant(
     return output, cost
 
 
-def which_tool(args: str) -> BaseModel:
-    adapter = TypeAdapter[
-        Confirmation
-        | BashCommand
-        | BashInteraction
-        | ResetShell
-        | Writefile
-        | FileEditFindReplace
-        | AIAssistant
-        | DoneFlag
-        | ReadImage
-    ](
-        Confirmation
-        | BashCommand
-        | BashInteraction
-        | ResetShell
-        | Writefile
-        | FileEditFindReplace
-        | AIAssistant
-        | DoneFlag
-        | ReadImage
-    )
+TOOLS = (
+    Confirmation
+    | BashCommand
+    | BashInteraction
+    | ResetShell
+    | Writefile
+    | CreateFileNew
+    | FileEditFindReplace
+    | FullFileEdit
+    | AIAssistant
+    | DoneFlag
+    | ReadImage
+)
+
+
+def which_tool(args: str) -> TOOLS:
+    adapter = TypeAdapter[TOOLS](TOOLS)
     return adapter.validate_python(json.loads(args))
 
 
+def which_tool_name(name: str) -> Type[TOOLS]:
+    if name == "Confirmation":
+        return Confirmation
+    elif name == "BashCommand":
+        return BashCommand
+    elif name == "BashInteraction":
+        return BashInteraction
+    elif name == "ResetShell":
+        return ResetShell
+    elif name == "Writefile":
+        return Writefile
+    elif name == "CreateFileNew":
+        return CreateFileNew
+    elif name == "FileEditFindReplace":
+        return FileEditFindReplace
+    elif name == "FullFileEdit":
+        return FullFileEdit
+    elif name == "AIAssistant":
+        return AIAssistant
+    elif name == "DoneFlag":
+        return DoneFlag
+    elif name == "ReadImage":
+        return ReadImage
+    else:
+        raise ValueError(f"Unknown tool name: {name}")
+
+
 def get_tool_output(
     args: dict[object, object]
     | Confirmation
@@ -518,7 +649,9 @@ def get_tool_output(
     | BashInteraction
     | ResetShell
     | Writefile
+    | CreateFileNew
     | FileEditFindReplace
+    | FullFileEdit
     | AIAssistant
     | DoneFlag
     | ReadImage,
@@ -534,7 +667,9 @@ def get_tool_output(
             | BashInteraction
             | ResetShell
             | Writefile
+            | CreateFileNew
             | FileEditFindReplace
+            | FullFileEdit
             | AIAssistant
             | DoneFlag
             | ReadImage
@@ -544,7 +679,9 @@ def get_tool_output(
             | BashInteraction
             | ResetShell
             | Writefile
+            | CreateFileNew
             | FileEditFindReplace
+            | FullFileEdit
             | AIAssistant
             | DoneFlag
             | ReadImage
@@ -561,10 +698,16 @@ def get_tool_output(
         output = execute_bash(enc, arg, max_tokens)
     elif isinstance(arg, Writefile):
         console.print("Calling write file tool")
-        output = write_file(arg), 0
+        output = write_file(arg, False), 0
+    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, FullFileEdit):
+        console.print("Calling full file edit tool")
+        output = do_diff_edit(arg), 0.0
     elif isinstance(arg, DoneFlag):
         console.print("Calling mark finish tool")
         output = mark_finish(arg), 0.0
@@ -593,44 +736,29 @@ curr_cost = 0.0
 
 
 class Mdata(BaseModel):
-    data: BashCommand | BashInteraction | Writefile | ResetShell | FileEditFindReplace
-
-
-execution_lock = threading.Lock()
-
-
-def execute_user_input() -> None:
-    while True:
-        discard_input()
-        user_input = input()
-        with execution_lock:
-            try:
-                console.log(
-                    execute_bash(
-                        default_enc,
-                        BashInteraction(
-                            send_ascii=[ord(x) for x in user_input] + [ord("\n")]
-                        ),
-                        max_tokens=None,
-                    )[0]
-                )
-            except Exception as e:
-                traceback.print_exc()
-                console.log(f"Error: {e}")
+    data: (
+        BashCommand
+        | BashInteraction
+        | Writefile
+        | CreateFileNew
+        | ResetShell
+        | FileEditFindReplace
+        | FullFileEdit
+    )
 
 
-async def register_client(server_url: str, client_uuid: str = "") -> None:
+def register_client(server_url: str, client_uuid: str = "") -> None:
     global default_enc, default_model, curr_cost
     # Generate a unique UUID for this client
     if not client_uuid:
         client_uuid = str(uuid.uuid4())
 
     # Create the WebSocket connection
-    async with websockets.connect(f"{server_url}/{client_uuid}") as websocket:
-        server_version = str(await websocket.recv())
+    with syncconnect(f"{server_url}/{client_uuid}") as websocket:
+        server_version = str(websocket.recv())
         print(f"Server version: {server_version}")
         client_version = importlib.metadata.version("wcgw")
-        await websocket.send(client_version)
+        websocket.send(client_version)
 
         print(
             f"Connected. Share this user id with the chatbot: {client_uuid} \nLink: https://chatgpt.com/g/g-Us0AAXkRh-wcgw-giving-shell-access"
@@ -638,24 +766,23 @@ async def register_client(server_url: str, client_uuid: str = "") -> None:
         try:
             while True:
                 # Wait to receive data from the server
-                message = await websocket.recv()
+                message = websocket.recv()
                 mdata = Mdata.model_validate_json(message)
-                with execution_lock:
-                    try:
-                        output, cost = get_tool_output(
-                            mdata.data, default_enc, 0.0, lambda x, y: ("", 0), None
-                        )
-                        curr_cost += cost
-                        print(f"{curr_cost=}")
-                    except Exception as e:
-                        output = f"GOT EXCEPTION while calling tool. Error: {e}"
-                        traceback.print_exc()
-                    assert isinstance(output, str)
-                    await websocket.send(output)
+                try:
+                    output, cost = get_tool_output(
+                        mdata.data, default_enc, 0.0, lambda x, y: ("", 0), None
+                    )
+                    curr_cost += cost
+                    print(f"{curr_cost=}")
+                except Exception as e:
+                    output = f"GOT EXCEPTION while calling tool. Error: {e}"
+                    traceback.print_exc()
+                assert isinstance(output, str)
+                websocket.send(output)
 
         except (websockets.ConnectionClosed, ConnectionError):
             print(f"Connection closed for UUID: {client_uuid}, retrying")
-            await register_client(server_url, client_uuid)
+            register_client(server_url, client_uuid)
 
 
 run = Typer(pretty_exceptions_show_locals=False, no_args_is_help=True)
@@ -672,13 +799,4 @@ def app(
         print(f"wcgw version: {version_}")
         exit()
 
-    thread1 = threading.Thread(target=execute_user_input)
-    thread2 = threading.Thread(
-        target=asyncio.run, args=(register_client(server_url, client_uuid or ""),)
-    )
-
-    thread1.start()
-    thread2.start()
-
-    thread1.join()
-    thread2.join()
+    register_client(server_url, client_uuid or "")

wcgw/relay/serve.py

@@ -17,7 +17,9 @@ from dotenv import load_dotenv
 from ..types_ import (
     BashCommand,
     BashInteraction,
+    CreateFileNew,
     FileEditFindReplace,
+    FullFileEdit,
     ResetShell,
     Writefile,
     Specials,
@@ -25,7 +27,15 @@ from ..types_ import (
 
 
 class Mdata(BaseModel):
-    data: BashCommand | BashInteraction | Writefile | ResetShell | FileEditFindReplace
+    data: (
+        BashCommand
+        | BashInteraction
+        | Writefile
+        | CreateFileNew
+        | ResetShell
+        | FileEditFindReplace
+        | FullFileEdit
+    )
     user_id: UUID
 
 
@@ -38,39 +48,6 @@ gpts: dict[UUID, Callable[[str], None]] = {}
 images: DefaultDict[UUID, dict[str, dict[str, Any]]] = DefaultDict(dict)
 
 
-@app.websocket("/register_serve_image/{uuid}")
-async def register_serve_image(websocket: WebSocket, uuid: UUID) -> None:
-    raise Exception("Disabled")
-    await websocket.accept()
-    received_data = await websocket.receive_json()
-    name = received_data["name"]
-    image_b64 = received_data["image_b64"]
-    image_bytes = base64.b64decode(image_b64)
-    images[uuid][name] = {
-        "content": image_bytes,
-        "media_type": received_data["media_type"],
-    }
-
-
-@app.get("/get_image/{uuid}/{name}")
-async def get_image(uuid: UUID, name: str) -> fastapi.responses.Response:
-    return fastapi.responses.Response(
-        content=images[uuid][name]["content"],
-        media_type=images[uuid][name]["media_type"],
-    )
-
-
-@app.websocket("/register/{uuid}")
-async def register_websocket_deprecated(websocket: WebSocket, uuid: UUID) -> None:
-    await websocket.accept()
-    await websocket.send_text(
-        "Outdated client used. Deprecated api is being used. Upgrade the wcgw app."
-    )
-    await websocket.close(
-        reason="This endpoint is deprecated. Please use /v1/register/{uuid}", code=1002
-    )
-
-
 CLIENT_VERSION_MINIMUM = "1.0.0"
 
 
@@ -116,25 +93,15 @@ async def register_websocket(websocket: WebSocket, uuid: UUID) -> None:
         print(f"Client {uuid} disconnected")
 
 
-@app.post("/write_file")
-async def write_file_deprecated(write_file_data: Writefile, user_id: UUID) -> Response:
-    return Response(
-        content="This version of the API is deprecated. Please upgrade your client.",
-        status_code=400,
-    )
-
-
-class WritefileWithUUID(Writefile):
+class CreateFileNewWithUUID(CreateFileNew):
     user_id: UUID
 
 
-@app.post("/v1/write_file")
-async def write_file(write_file_data: WritefileWithUUID) -> str:
+@app.post("/v1/create_file")
+async def create_file(write_file_data: CreateFileNewWithUUID) -> str:
     user_id = write_file_data.user_id
     if user_id not in clients:
-        raise fastapi.HTTPException(
-            status_code=404, detail="User with the provided id not found"
-        )
+        return "Failure: id not found, ask the user to check it."
 
     results: Optional[str] = None
 
@@ -155,19 +122,17 @@ async def write_file(write_file_data: WritefileWithUUID) -> str:
     raise fastapi.HTTPException(status_code=500, detail="Timeout error")
 
 
-class FileEditFindReplaceWithUUID(FileEditFindReplace):
+class FullFileEditWithUUID(FullFileEdit):
     user_id: UUID
 
 
-@app.post("/v1/file_edit_find_replace")
+@app.post("/v1/full_file_edit")
 async def file_edit_find_replace(
-    file_edit_find_replace: FileEditFindReplaceWithUUID,
+    file_edit_find_replace: FullFileEditWithUUID,
 ) -> str:
     user_id = file_edit_find_replace.user_id
     if user_id not in clients:
-        raise fastapi.HTTPException(
-            status_code=404, detail="User with the provided id not found"
-        )
+        return "Failure: id not found, ask the user to check it."
 
     results: Optional[str] = None
 
@@ -201,9 +166,7 @@ class ResetShellWithUUID(ResetShell):
 async def reset_shell(reset_shell: ResetShellWithUUID) -> str:
     user_id = reset_shell.user_id
     if user_id not in clients:
-        raise fastapi.HTTPException(
-            status_code=404, detail="User with the provided id not found"
-        )
+        return "Failure: id not found, ask the user to check it."
 
     results: Optional[str] = None
 
@@ -224,14 +187,6 @@ async def reset_shell(reset_shell: ResetShellWithUUID) -> str:
     raise fastapi.HTTPException(status_code=500, detail="Timeout error")
 
 
-@app.post("/execute_bash")
-async def execute_bash_deprecated(excute_bash_data: Any, user_id: UUID) -> Response:
-    return Response(
-        content="This version of the API is deprecated. Please upgrade your client.",
-        status_code=400,
-    )
-
-
 class CommandWithUUID(BaseModel):
     command: str
     user_id: UUID
@@ -241,9 +196,7 @@ class CommandWithUUID(BaseModel):
 async def bash_command(command: CommandWithUUID) -> str:
     user_id = command.user_id
     if user_id not in clients:
-        raise fastapi.HTTPException(
-            status_code=404, detail="User with the provided id not found"
-        )
+        return "Failure: id not found, ask the user to check it."
 
     results: Optional[str] = None
 
@@ -274,9 +227,7 @@ class BashInteractionWithUUID(BashInteraction):
 async def bash_interaction(bash_interaction: BashInteractionWithUUID) -> str:
     user_id = bash_interaction.user_id
     if user_id not in clients:
-        raise fastapi.HTTPException(
-            status_code=404, detail="User with the provided id not found"
-        )
+        return "Failure: id not found, ask the user to check it."
 
     results: Optional[str] = None
 

wcgw/types_.py

@@ -16,6 +16,22 @@ class BashInteraction(BaseModel):
     send_specials: Optional[Sequence[Specials]] = None
     send_ascii: Optional[Sequence[int]] = None
 
+    def model_post_init(self, __context: object) -> None:
+        # Ensure only one of the fields is set
+        if (
+            sum(
+                [
+                    int(bool(self.send_text)),
+                    int(bool(self.send_specials)),
+                    int(bool(self.send_ascii)),
+                ]
+            )
+            != 1
+        ):
+            raise ValueError(
+                "Exactly one of 'send_text', 'send_specials', or 'send_ascii' must be set"
+            )
+
 
 class ReadImage(BaseModel):
     file_path: str
@@ -27,6 +43,11 @@ class Writefile(BaseModel):
     file_content: str
 
 
+class CreateFileNew(BaseModel):
+    file_path: str
+    file_content: str
+
+
 class FileEditFindReplace(BaseModel):
     file_path: str
     find_lines: str
@@ -35,3 +56,8 @@ class FileEditFindReplace(BaseModel):
 
 class ResetShell(BaseModel):
     should_reset: Literal[True] = True
+
+
+class FullFileEdit(BaseModel):
+    file_path: str
+    file_edit_using_searh_replace_blocks: str