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

wcgw/__init__.py

@@ -1,2 +1,2 @@
-from .basic import app, loop
-from .tools import run as listen
+from .client.basic import app, loop
+from .client.tools import run as listen

wcgw/client/__main__.py

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

wcgw/{basic.py → client/basic.py}

@@ -20,16 +20,15 @@ import petname  # type: ignore[import-untyped]
 from typer import Typer
 import uuid
 
+from ..types_ import BashCommand, BashInteraction, ReadImage, Writefile, ResetShell
+
 from .common import Models, discard_input
 from .common import CostData, History
 from .openai_utils import get_input_cost, get_output_cost
-from .tools import ExecuteBash, ReadImage, ImageData
+from .tools import ImageData
 
 from .tools import (
-    BASH_CLF_OUTPUT,
-    Confirmation,
     DoneFlag,
-    Writefile,
     get_tool_output,
     SHELL,
     start_shell,
@@ -156,26 +155,32 @@ def loop(
 
     tools = [
         openai.pydantic_function_tool(
-            ExecuteBash,
+            BashCommand,
             description="""
-- Execute a bash script. This is stateful (beware with subsequent calls).
-- Execute commands using `execute_command` attribute.
+- Execute a bash command. This is stateful (beware with subsequent calls).
 - Do not use interactive commands like nano. Prefer writing simpler commands.
-- Last line will always be `(exit <int code>)` except if
-- The last line is `(pending)` if the program is still running or waiting for your input. You can then send input using `send_ascii` attributes. You get status by sending new line `send_ascii: ["Enter"]` or `send_ascii: [10]`.
-- Optionally the last line is `(won't exit)` in which case you need to kill the process if you want to run a new command.
+- Status of the command and the current working directory will always be returned at the end.
 - Optionally `exit shell has restarted` is the output, in which case environment resets, you can run fresh commands.
 - 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.
 """,
+        ),
+        openai.pydantic_function_tool(
+            BashInteraction,
+            description="""
+- Interact with running program using this tool.""",
         ),
         openai.pydantic_function_tool(
             Writefile,
-            description="Write content to a file. Provide file path and content. Use this instead of ExecuteBash for writing files.",
+            description="Write content to a file. Provide file path and content. Use this instead of BashCommand for writing files.",
         ),
         openai.pydantic_function_tool(
             ReadImage, description="Read an image from the shell."
         ),
+        openai.pydantic_function_tool(
+            ResetShell,
+            description="Resets the shell. Use only if all interrupts and prompt reset attempts have failed repeatedly.",
+        ),
     ]
     uname_sysname = os.uname().sysname
     uname_machine = os.uname().machine

wcgw/{tools.py → client/tools.py}

@@ -2,8 +2,10 @@ import asyncio
 import base64
 import json
 import mimetypes
+import re
 import sys
 import threading
+import importlib.metadata
 import traceback
 from typing import (
     Callable,
@@ -11,12 +13,12 @@ from typing import (
     NewType,
     Optional,
     ParamSpec,
-    Sequence,
     TypeVar,
     TypedDict,
 )
 import uuid
 from pydantic import BaseModel, TypeAdapter
+import typer
 from websockets.sync.client import connect as syncconnect
 
 import os
@@ -38,6 +40,14 @@ from openai.types.chat import (
     ChatCompletionMessage,
     ParsedChatCompletionMessage,
 )
+from nltk.metrics.distance import edit_distance
+from ..types_ import FileEditFindReplace, ResetShell, Writefile
+
+from ..types_ import BashCommand
+
+from ..types_ import BashInteraction
+
+from ..types_ import ReadImage
 
 from .common import CostData, Models, discard_input
 
@@ -72,22 +82,20 @@ def ask_confirmation(prompt: Confirmation) -> str:
     return "Yes" if response.lower() == "y" else "No"
 
 
-class Writefile(BaseModel):
-    file_path: str
-    file_content: str
+PROMPT = "#@@"
 
 
-def start_shell() -> pexpect.spawn:
+def start_shell() -> pexpect.spawn:  # type: ignore
     SHELL = pexpect.spawn(
         "/bin/bash --noprofile --norc",
-        env={**os.environ, **{"PS1": "#@@"}},  # type: ignore[arg-type]
+        env={**os.environ, **{"PS1": PROMPT}},  # type: ignore[arg-type]
         echo=False,
         encoding="utf-8",
         timeout=TIMEOUT,
     )
-    SHELL.expect("#@@")
+    SHELL.expect(PROMPT)
     SHELL.sendline("stty -icanon -echo")
-    SHELL.expect("#@@")
+    SHELL.expect(PROMPT)
     return SHELL
 
 
@@ -103,16 +111,22 @@ def _is_int(mystr: str) -> bool:
 
 
 def _get_exit_code() -> int:
+    if PROMPT != "#@@":
+        return 0
     # First reset the prompt in case venv was sourced or other reasons.
-    SHELL.sendline('export PS1="#@@"')
-    SHELL.expect("#@@")
+    SHELL.sendline(f"export PS1={PROMPT}")
+    SHELL.expect(PROMPT)
     # Reset echo also if it was enabled
     SHELL.sendline("stty -icanon -echo")
-    SHELL.expect("#@@")
+    SHELL.expect(PROMPT)
     SHELL.sendline("echo $?")
     before = ""
     while not _is_int(before):  # Consume all previous output
-        SHELL.expect("#@@")
+        try:
+            SHELL.expect(PROMPT)
+        except pexpect.TIMEOUT:
+            print(f"Couldn't get exit code, before: {before}")
+            raise
         assert isinstance(SHELL.before, str)
         # Render because there could be some anscii escape sequences still set like in google colab env
         before = render_terminal_output(SHELL.before).strip()
@@ -123,51 +137,110 @@ def _get_exit_code() -> int:
         raise ValueError(f"Malformed output: {before}")
 
 
-Specials = Literal["Key-up", "Key-down", "Key-left", "Key-right", "Enter", "Ctrl-c"]
+BASH_CLF_OUTPUT = Literal["repl", "pending"]
+BASH_STATE: BASH_CLF_OUTPUT = "repl"
+CWD = os.getcwd()
+
+
+def reset_shell() -> str:
+    global SHELL, BASH_STATE, CWD
+    SHELL.close(True)
+    SHELL = start_shell()
+    BASH_STATE = "repl"
+    CWD = os.getcwd()
+    return "Reset successful" + get_status()
+
+
+WAITING_INPUT_MESSAGE = """A command is already running. NOTE: You can't run multiple shell sessions, likely a previous program hasn't exited. 
+1. Get its output using `send_ascii: [10] or send_specials: ["Enter"]`
+2. Use `send_ascii` or `send_specials` to give inputs to the running program, don't use `BashCommand` OR
+3. kill the previous program by sending ctrl+c first using `send_ascii` or `send_specials`
+4. Send the process in background using `send_specials: ["Ctrl-z"]` followed by BashCommand: `bg`
+"""
+
+
+def update_repl_prompt(command: str) -> bool:
+    global PROMPT
+    if re.match(r"^wcgw_update_prompt\(\)$", command.strip()):
+        SHELL.sendintr()
+        index = SHELL.expect([PROMPT, pexpect.TIMEOUT], timeout=0.2)
+        if index == 0:
+            return False
+        before = SHELL.before or ""
+        assert before, "Something went wrong updating repl prompt"
+        PROMPT = before.split("\n")[-1].strip()
+        # Escape all regex
+        PROMPT = re.escape(PROMPT)
+        print(f"Trying to update prompt to: {PROMPT.encode()!r}")
+        index = 0
+        while index == 0:
+            # Consume all REPL prompts till now
+            index = SHELL.expect([PROMPT, pexpect.TIMEOUT], timeout=0.2)
+        print(f"Prompt updated to: {PROMPT}")
+        return True
+    return False
 
 
-class ExecuteBash(BaseModel):
-    execute_command: Optional[str] = None
-    send_ascii: Optional[Sequence[int | Specials]] = None
+def get_cwd() -> str:
+    SHELL.sendline("pwd")
+    SHELL.expect(PROMPT)
+    assert isinstance(SHELL.before, str)
+    current_dir = render_terminal_output(SHELL.before).strip()
+    return current_dir
 
 
-BASH_CLF_OUTPUT = Literal["running", "waiting_for_input", "wont_exit"]
-BASH_STATE: BASH_CLF_OUTPUT = "running"
+def get_status() -> str:
+    global CWD
+    exit_code: Optional[int] = None
 
+    status = "\n\n---\n\n"
+    if BASH_STATE == "pending":
+        status += "status = still running\n"
+        status += "cwd = " + CWD + "\n"
+    else:
+        exit_code = _get_exit_code()
+        status += f"status = exited with code {exit_code}\n"
+        CWD = get_cwd()
+        status += "cwd = " + CWD + "\n"
 
-WAITING_INPUT_MESSAGE = """A command is already running waiting for input. NOTE: You can't run multiple shell sessions, likely a previous program hasn't exited. 
-1. Get its output using `send_ascii: [10]`
-2. Use `send_ascii` to give inputs to the running program, don't use `execute_command` OR
-3. kill the previous program by sending ctrl+c first using `send_ascii`"""
+    return status.rstrip()
 
 
 def execute_bash(
-    enc: tiktoken.Encoding, bash_arg: ExecuteBash, max_tokens: Optional[int]
+    enc: tiktoken.Encoding,
+    bash_arg: BashCommand | BashInteraction,
+    max_tokens: Optional[int],
 ) -> tuple[str, float]:
-    global SHELL, BASH_STATE
+    global SHELL, BASH_STATE, CWD
     try:
-        if bash_arg.execute_command:
-            if BASH_STATE == "waiting_for_input":
-                raise ValueError(WAITING_INPUT_MESSAGE)
-            elif BASH_STATE == "wont_exit":
-                raise ValueError(
-                    """A command is already running that hasn't exited. NOTE: You can't run multiple shell sessions, likely a previous program is in infinite loop.
-                    Kill the previous program by sending ctrl+c first using `send_ascii`"""
+        is_interrupt = False
+        if isinstance(bash_arg, BashCommand):
+            updated_repl_mode = update_repl_prompt(bash_arg.command)
+            if updated_repl_mode:
+                BASH_STATE = "repl"
+                response = (
+                    "Prompt updated, you can execute REPL lines using BashCommand now"
                 )
-            command = bash_arg.execute_command.strip()
+                console.print(response)
+                return (
+                    response,
+                    0,
+                )
+
+            console.print(f"$ {bash_arg.command}")
+            if BASH_STATE == "pending":
+                raise ValueError(WAITING_INPUT_MESSAGE)
+            command = bash_arg.command.strip()
 
             if "\n" in command:
                 raise ValueError(
                     "Command should not contain newline character in middle. Run only one command at a time."
                 )
 
-            console.print(f"$ {command}")
             SHELL.sendline(command)
-        elif bash_arg.send_ascii:
-            console.print(f"Sending ASCII sequence: {bash_arg.send_ascii}")
-            for char in bash_arg.send_ascii:
-                if isinstance(char, int):
-                    SHELL.send(chr(char))
+        elif bash_arg.send_specials:
+            console.print(f"Sending special sequence: {bash_arg.send_specials}")
+            for char in bash_arg.send_specials:
                 if char == "Key-up":
                     SHELL.send("\033[A")
                 elif char == "Key-down":
@@ -180,26 +253,53 @@ def execute_bash(
                     SHELL.send("\n")
                 elif char == "Ctrl-c":
                     SHELL.sendintr()
+                    is_interrupt = True
+                elif char == "Ctrl-d":
+                    SHELL.sendintr()
+                    is_interrupt = True
+                elif char == "Ctrl-z":
+                    SHELL.send("\x1a")
+                else:
+                    raise Exception(f"Unknown special character: {char}")
+        elif bash_arg.send_ascii:
+            console.print(f"Sending ASCII sequence: {bash_arg.send_ascii}")
+            for ascii_char in bash_arg.send_ascii:
+                SHELL.send(chr(ascii_char))
+                if ascii_char == 3:
+                    is_interrupt = True
         else:
-            raise Exception("Nothing to send")
-        BASH_STATE = "running"
+            if bash_arg.send_text is None:
+                return (
+                    "Failure: at least one of send_text, send_specials or send_ascii should be provided",
+                    0.0,
+                )
+
+            updated_repl_mode = update_repl_prompt(bash_arg.send_text)
+            if updated_repl_mode:
+                BASH_STATE = "repl"
+                response = (
+                    "Prompt updated, you can execute REPL lines using BashCommand now"
+                )
+                console.print(response)
+                return (
+                    response,
+                    0,
+                )
+            console.print(f"Interact text: {bash_arg.send_text}")
+            SHELL.sendline(bash_arg.send_text)
+
+        BASH_STATE = "repl"
 
     except KeyboardInterrupt:
-        SHELL.close(True)
-        SHELL = start_shell()
-        raise
+        SHELL.sendintr()
+        SHELL.expect(PROMPT)
+        return "---\n\nFailure: user interrupted the execution", 0.0
 
     wait = 5
-    index = SHELL.expect(["#@@", pexpect.TIMEOUT], timeout=wait)
-    running = ""
-    while index == 1:
-        if wait > TIMEOUT:
-            raise TimeoutError("Timeout while waiting for shell prompt")
-
-        BASH_STATE = "waiting_for_input"
+    index = SHELL.expect([PROMPT, pexpect.TIMEOUT], timeout=wait)
+    if index == 1:
+        BASH_STATE = "pending"
         text = SHELL.before or ""
-        print(text[len(running) :])
-        running = text
 
         text = render_terminal_output(text)
         tokens = enc.encode(text)
@@ -207,8 +307,25 @@ def execute_bash(
         if max_tokens and len(tokens) >= max_tokens:
             text = "...(truncated)\n" + enc.decode(tokens[-(max_tokens - 1) :])
 
-        last_line = "(pending)"
-        return text + f"\n{last_line}", 0
+        if is_interrupt:
+            text = (
+                text
+                + """---
+----
+Failure interrupting.
+If any REPL session was previously running or if bashrc was sourced, or if there is issue to other REPL related reasons:
+    Run BashCommand: "wcgw_update_prompt()" to reset the PS1 prompt.
+Otherwise, you may want to try Ctrl-c again or program specific exit interactive commands.
+"""
+            )
+
+        exit_status = get_status()
+        text += exit_status
+
+        return text, 0
+
+    if is_interrupt:
+        return "Interrupt successful", 0.0
 
     assert isinstance(SHELL.before, str)
     output = render_terminal_output(SHELL.before)
@@ -218,9 +335,8 @@ def execute_bash(
         output = "...(truncated)\n" + enc.decode(tokens[-(max_tokens - 1) :])
 
     try:
-        exit_code = _get_exit_code()
-        output += f"\n(exit {exit_code})"
-
+        exit_status = get_status()
+        output += exit_status
     except ValueError as e:
         console.print(output)
         traceback.print_exc()
@@ -232,11 +348,6 @@ def execute_bash(
     return output, 0
 
 
-class ReadImage(BaseModel):
-    file_path: str
-    type: Literal["ReadImage"] = "ReadImage"
-
-
 def serve_image_in_bg(file_path: str, client_uuid: str, name: str) -> None:
     if not client_uuid:
         client_uuid = str(uuid.uuid4())
@@ -269,25 +380,17 @@ T = TypeVar("T")
 def ensure_no_previous_output(func: Callable[Param, T]) -> Callable[Param, T]:
     def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> T:
         global BASH_STATE
-        if BASH_STATE == "waiting_for_input":
+        if BASH_STATE == "pending":
             raise ValueError(WAITING_INPUT_MESSAGE)
-        elif BASH_STATE == "wont_exit":
-            raise ValueError(
-                "A command is already running that hasn't exited. NOTE: You can't run multiple shell sessions, likely the previous program is in infinite loop. Please kill the previous program by sending ctrl+c first."
-            )
+
         return func(*args, **kwargs)
 
     return wrapper
 
 
-@ensure_no_previous_output
 def read_image_from_shell(file_path: str) -> ImageData:
     if not os.path.isabs(file_path):
-        SHELL.sendline("pwd")
-        SHELL.expect("#@@")
-        assert isinstance(SHELL.before, str)
-        current_dir = render_terminal_output(SHELL.before).strip()
-        file_path = os.path.join(current_dir, 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")
@@ -299,22 +402,65 @@ def read_image_from_shell(file_path: str) -> ImageData:
         return ImageData(dataurl=f"data:{image_type};base64,{image_b64}")
 
 
-@ensure_no_previous_output
 def write_file(writefile: Writefile) -> str:
     if not os.path.isabs(writefile.file_path):
-        SHELL.sendline("pwd")
-        SHELL.expect("#@@")
-        assert isinstance(SHELL.before, str)
-        current_dir = render_terminal_output(SHELL.before).strip()
-        return f"Failure: Use absolute path only. FYI current working directory is '{current_dir}'"
-    os.makedirs(os.path.dirname(writefile.file_path), exist_ok=True)
+        path_ = os.path.join(CWD, writefile.file_path)
+    else:
+        path_ = writefile.file_path
     try:
-        with open(writefile.file_path, "w") as f:
+        with open(path_, "w") as f:
             f.write(writefile.file_content)
     except OSError as e:
-        console.print(f"Error: {e}", style="red")
         return f"Error: {e}"
-    console.print(f"File written to {writefile.file_path}")
+    console.print(f"File written to {path_}")
+    return "Success"
+
+
+def find_least_edit_distance_substring(content: str, find_str: str) -> str:
+    content_lines = content.split("\n")
+    find_lines = find_str.split("\n")
+    # 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):
+        edit_distance_sum = 0
+        for j in range(len(find_lines)):
+            edit_distance_sum += edit_distance(content_lines[i + j], 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)
+
+
+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)
+    else:
+        path_ = file_edit.file_path
+
+    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")
+    )
+    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 = 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}"
+    console.print(f"File written to {path_}")
     return "Success"
 
 
@@ -342,16 +488,37 @@ def take_help_of_ai_assistant(
 
 def which_tool(args: str) -> BaseModel:
     adapter = TypeAdapter[
-        Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag | ReadImage
-    ](Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag | ReadImage)
+        Confirmation
+        | BashCommand
+        | BashInteraction
+        | ResetShell
+        | Writefile
+        | FileEditFindReplace
+        | AIAssistant
+        | DoneFlag
+        | ReadImage
+    ](
+        Confirmation
+        | BashCommand
+        | BashInteraction
+        | ResetShell
+        | Writefile
+        | FileEditFindReplace
+        | AIAssistant
+        | DoneFlag
+        | ReadImage
+    )
     return adapter.validate_python(json.loads(args))
 
 
 def get_tool_output(
     args: dict[object, object]
     | Confirmation
-    | ExecuteBash
+    | BashCommand
+    | BashInteraction
+    | ResetShell
     | Writefile
+    | FileEditFindReplace
     | AIAssistant
     | DoneFlag
     | ReadImage,
@@ -362,8 +529,26 @@ def get_tool_output(
 ) -> tuple[str | ImageData | DoneFlag, float]:
     if isinstance(args, dict):
         adapter = TypeAdapter[
-            Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag | ReadImage
-        ](Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag | ReadImage)
+            Confirmation
+            | BashCommand
+            | BashInteraction
+            | ResetShell
+            | Writefile
+            | FileEditFindReplace
+            | AIAssistant
+            | DoneFlag
+            | ReadImage
+        ](
+            Confirmation
+            | BashCommand
+            | BashInteraction
+            | ResetShell
+            | Writefile
+            | FileEditFindReplace
+            | AIAssistant
+            | DoneFlag
+            | ReadImage
+        )
         arg = adapter.validate_python(args)
     else:
         arg = args
@@ -371,12 +556,15 @@ def get_tool_output(
     if isinstance(arg, Confirmation):
         console.print("Calling ask confirmation tool")
         output = ask_confirmation(arg), 0.0
-    elif isinstance(arg, ExecuteBash):
+    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), 0
+    elif isinstance(arg, FileEditFindReplace):
+        console.print("Calling file edit tool")
+        output = file_edit(arg), 0.0
     elif isinstance(arg, DoneFlag):
         console.print("Calling mark finish tool")
         output = mark_finish(arg), 0.0
@@ -386,6 +574,9 @@ def get_tool_output(
     elif isinstance(arg, ReadImage):
         console.print("Calling read image tool")
         output = read_image_from_shell(arg.file_path), 0.0
+    elif isinstance(arg, ResetShell):
+        console.print("Calling reset shell tool")
+        output = reset_shell(), 0.0
     else:
         raise ValueError(f"Unknown tool: {arg}")
 
@@ -395,44 +586,6 @@ def get_tool_output(
 
 History = list[ChatCompletionMessageParam]
 
-
-def get_is_waiting_user_input(
-    model: Models, cost_data: CostData
-) -> Callable[[str], tuple[BASH_CLF_OUTPUT, float]]:
-    enc = tiktoken.encoding_for_model(model if not model.startswith("o1") else "gpt-4o")
-    system_prompt = """You need to classify if a bash program is waiting for user input based on its stdout, or if it won't exit. You'll be given the output of any program.
-    Return `waiting_for_input` if the program is waiting for INTERACTIVE input only, Return 'running' if it's waiting for external resources or just waiting to finish.
-    Return `wont_exit` if the program won't exit, for example if it's a server.
-    Return `running` otherwise.
-    """
-    history: History = [{"role": "system", "content": system_prompt}]
-    client = OpenAI()
-
-    class ExpectedOutput(BaseModel):
-        output_classified: BASH_CLF_OUTPUT
-
-    def is_waiting_user_input(output: str) -> tuple[BASH_CLF_OUTPUT, float]:
-        # Send only last 30 lines
-        output = "\n".join(output.split("\n")[-30:])
-        # Send only max last 200 tokens
-        output = enc.decode(enc.encode(output)[-200:])
-
-        history.append({"role": "user", "content": output})
-        response = client.beta.chat.completions.parse(
-            model=model, messages=history, response_format=ExpectedOutput
-        )
-        parsed = response.choices[0].message.parsed
-        if parsed is None:
-            raise ValueError("No parsed output")
-        cost = (
-            get_input_cost(cost_data, enc, history)[0]
-            + get_output_cost(cost_data, enc, response.choices[0].message)[0]
-        )
-        return parsed.output_classified, cost
-
-    return is_waiting_user_input
-
-
 default_enc = tiktoken.encoding_for_model("gpt-4o")
 default_model: Models = "gpt-4o-2024-08-06"
 default_cost = CostData(cost_per_1m_input_tokens=0.15, cost_per_1m_output_tokens=0.6)
@@ -440,7 +593,7 @@ curr_cost = 0.0
 
 
 class Mdata(BaseModel):
-    data: ExecuteBash | Writefile
+    data: BashCommand | BashInteraction | Writefile | ResetShell | FileEditFindReplace
 
 
 execution_lock = threading.Lock()
@@ -455,7 +608,7 @@ def execute_user_input() -> None:
                 console.log(
                     execute_bash(
                         default_enc,
-                        ExecuteBash(
+                        BashInteraction(
                             send_ascii=[ord(x) for x in user_input] + [ord("\n")]
                         ),
                         max_tokens=None,
@@ -474,6 +627,11 @@ async def register_client(server_url: str, client_uuid: str = "") -> None:
 
     # Create the WebSocket connection
     async with websockets.connect(f"{server_url}/{client_uuid}") as websocket:
+        server_version = str(await websocket.recv())
+        print(f"Server version: {server_version}")
+        client_version = importlib.metadata.version("wcgw")
+        await 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"
         )
@@ -505,8 +663,15 @@ run = Typer(pretty_exceptions_show_locals=False, no_args_is_help=True)
 
 @run.command()
 def app(
-    server_url: str = "wss://wcgw.arcfu.com/register", client_uuid: Optional[str] = None
+    server_url: str = "wss://wcgw.arcfu.com/v1/register",
+    client_uuid: Optional[str] = None,
+    version: bool = typer.Option(False, "--version", "-v"),
 ) -> None:
+    if version:
+        version_ = importlib.metadata.version("wcgw")
+        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 ""),)

wcgw/relay/serve.py

@@ -0,0 +1,326 @@
+import asyncio
+import base64
+from importlib import metadata
+import semantic_version  # type: ignore[import-untyped]
+import threading
+import time
+from typing import Any, Callable, Coroutine, DefaultDict, Literal, Optional, Sequence
+from uuid import UUID
+import fastapi
+from fastapi import Response, WebSocket, WebSocketDisconnect
+from pydantic import BaseModel
+import uvicorn
+from fastapi.staticfiles import StaticFiles
+
+from dotenv import load_dotenv
+
+from ..types_ import (
+    BashCommand,
+    BashInteraction,
+    FileEditFindReplace,
+    ResetShell,
+    Writefile,
+    Specials,
+)
+
+
+class Mdata(BaseModel):
+    data: BashCommand | BashInteraction | Writefile | ResetShell | FileEditFindReplace
+    user_id: UUID
+
+
+app = fastapi.FastAPI()
+
+clients: dict[UUID, Callable[[Mdata], Coroutine[None, None, None]]] = {}
+websockets: dict[UUID, WebSocket] = {}
+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"
+
+
+@app.websocket("/v1/register/{uuid}")
+async def register_websocket(websocket: WebSocket, uuid: UUID) -> None:
+    await websocket.accept()
+
+    # send server version
+    version = metadata.version("wcgw")
+    await websocket.send_text(version)
+
+    # 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)
+    if sem_version_client < sem_version_server:
+        await websocket.send_text(
+            f"Client version {client_version} is outdated. Please upgrade to {CLIENT_VERSION_MINIMUM} or higher."
+        )
+        await websocket.close(
+            reason="Client version outdated. Please upgrade to the latest version.",
+            code=1002,
+        )
+        return
+
+    # Register the callback for this client UUID
+    async def send_data_callback(data: Mdata) -> None:
+        await websocket.send_text(data.model_dump_json())
+
+    clients[uuid] = send_data_callback
+    websockets[uuid] = websocket
+
+    try:
+        while True:
+            received_data = await websocket.receive_text()
+            if uuid not in gpts:
+                raise fastapi.HTTPException(status_code=400, detail="No call made")
+            gpts[uuid](received_data)
+    except WebSocketDisconnect:
+        # Remove the client if the WebSocket is disconnected
+        del clients[uuid]
+        del websockets[uuid]
+        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):
+    user_id: UUID
+
+
+@app.post("/v1/write_file")
+async def write_file(write_file_data: WritefileWithUUID) -> 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"
+        )
+
+    results: Optional[str] = None
+
+    def put_results(result: str) -> None:
+        nonlocal results
+        results = result
+
+    gpts[user_id] = put_results
+
+    await clients[user_id](Mdata(data=write_file_data, user_id=user_id))
+
+    start_time = time.time()
+    while time.time() - start_time < 30:
+        if results is not None:
+            return results
+        await asyncio.sleep(0.1)
+
+    raise fastapi.HTTPException(status_code=500, detail="Timeout error")
+
+
+class FileEditFindReplaceWithUUID(FileEditFindReplace):
+    user_id: UUID
+
+
+@app.post("/v1/file_edit_find_replace")
+async def file_edit_find_replace(
+    file_edit_find_replace: FileEditFindReplaceWithUUID,
+) -> 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"
+        )
+
+    results: Optional[str] = None
+
+    def put_results(result: str) -> None:
+        nonlocal results
+        results = result
+
+    gpts[user_id] = put_results
+
+    await clients[user_id](
+        Mdata(
+            data=file_edit_find_replace,
+            user_id=user_id,
+        )
+    )
+
+    start_time = time.time()
+    while time.time() - start_time < 30:
+        if results is not None:
+            return results
+        await asyncio.sleep(0.1)
+
+    raise fastapi.HTTPException(status_code=500, detail="Timeout error")
+
+
+class ResetShellWithUUID(ResetShell):
+    user_id: UUID
+
+
+@app.post("/v1/reset_shell")
+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"
+        )
+
+    results: Optional[str] = None
+
+    def put_results(result: str) -> None:
+        nonlocal results
+        results = result
+
+    gpts[user_id] = put_results
+
+    await clients[user_id](Mdata(data=reset_shell, user_id=user_id))
+
+    start_time = time.time()
+    while time.time() - start_time < 30:
+        if results is not None:
+            return results
+        await asyncio.sleep(0.1)
+
+    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
+
+
+@app.post("/v1/bash_command")
+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"
+        )
+
+    results: Optional[str] = None
+
+    def put_results(result: str) -> None:
+        nonlocal results
+        results = result
+
+    gpts[user_id] = put_results
+
+    await clients[user_id](
+        Mdata(data=BashCommand(command=command.command), user_id=user_id)
+    )
+
+    start_time = time.time()
+    while time.time() - start_time < 30:
+        if results is not None:
+            return results
+        await asyncio.sleep(0.1)
+
+    raise fastapi.HTTPException(status_code=500, detail="Timeout error")
+
+
+class BashInteractionWithUUID(BashInteraction):
+    user_id: UUID
+
+
+@app.post("/v1/bash_interaction")
+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"
+        )
+
+    results: Optional[str] = None
+
+    def put_results(result: str) -> None:
+        nonlocal results
+        results = result
+
+    gpts[user_id] = put_results
+
+    await clients[user_id](
+        Mdata(
+            data=bash_interaction,
+            user_id=user_id,
+        )
+    )
+
+    start_time = time.time()
+    while time.time() - start_time < 30:
+        if results is not None:
+            return results
+        await asyncio.sleep(0.1)
+
+    raise fastapi.HTTPException(status_code=500, detail="Timeout error")
+
+
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+
+def run() -> None:
+    load_dotenv()
+
+    uvicorn_thread = threading.Thread(
+        target=uvicorn.run,
+        args=(app,),
+        kwargs={
+            "host": "0.0.0.0",
+            "port": 8000,
+            "log_level": "info",
+            "access_log": True,
+        },
+    )
+    uvicorn_thread.start()
+    uvicorn_thread.join()
+
+
+if __name__ == "__main__":
+    run()

wcgw/relay/static/privacy.txt

@@ -0,0 +1,7 @@
+Privacy Policy
+I do not collect, store, or share any personal data.
+The data from your terminal is not stored anywhere and it's not logged or collected in any form.
+There is a relay webserver for connecting your terminal to chatgpt the source code for which is open at https://github.com/rusiaaman/wcgw/tree/main/src/relay that you can run on your own.
+Other than the relay webserver there is no further involvement of my servers or services.
+Feel free to me contact at info@arcfu.com for questions on privacy or anything else.
+

wcgw/types_.py

@@ -0,0 +1,37 @@
+from typing import Literal, Optional, Sequence
+from pydantic import BaseModel
+
+
+class BashCommand(BaseModel):
+    command: str
+
+
+Specials = Literal[
+    "Key-up", "Key-down", "Key-left", "Key-right", "Enter", "Ctrl-c", "Ctrl-d", "Ctrl-z"
+]
+
+
+class BashInteraction(BaseModel):
+    send_text: Optional[str] = None
+    send_specials: Optional[Sequence[Specials]] = None
+    send_ascii: Optional[Sequence[int]] = None
+
+
+class ReadImage(BaseModel):
+    file_path: str
+    type: Literal["ReadImage"] = "ReadImage"
+
+
+class Writefile(BaseModel):
+    file_path: str
+    file_content: str
+
+
+class FileEditFindReplace(BaseModel):
+    file_path: str
+    find_lines: str
+    replace_with_lines: str
+
+
+class ResetShell(BaseModel):
+    should_reset: Literal[True] = True

{wcgw-0.1.2.dist-info → wcgw-1.0.0.dist-info}/METADATA

@@ -1,12 +1,13 @@
 Metadata-Version: 2.3
 Name: wcgw
-Version: 0.1.2
+Version: 1.0.0
 Summary: What could go wrong giving full shell access to chatgpt?
 Project-URL: Homepage, https://github.com/rusiaaman/wcgw
 Author-email: Aman Rusia <gapypi@arcfu.com>
 Requires-Python: <3.13,>=3.10
 Requires-Dist: fastapi>=0.115.0
 Requires-Dist: mypy>=1.11.2
+Requires-Dist: nltk>=3.9.1
 Requires-Dist: openai>=1.46.0
 Requires-Dist: petname>=2.6
 Requires-Dist: pexpect>=4.9.0
@@ -14,6 +15,7 @@ Requires-Dist: pydantic>=2.9.2
 Requires-Dist: pyte>=0.8.2
 Requires-Dist: python-dotenv>=1.0.1
 Requires-Dist: rich>=13.8.1
+Requires-Dist: semantic-version>=2.10.0
 Requires-Dist: shell>=1.0.1
 Requires-Dist: tiktoken==0.7.0
 Requires-Dist: toml>=0.10.2
@@ -32,7 +34,8 @@ A custom gpt on chatgpt web app to interact with your local shell.
 ### 🚀 Highlights
 - ⚡ **Full Shell Access**: No restrictions, complete control.
 - ⚡ **Create, Execute, Iterate**: Ask the gpt to keep running compiler checks till all errors are fixed, or ask it to keep checking for the status of a long running command till it's done.
-- ⚡ **Interactive Command Handling**: [beta] Supports interactive commands using arrow keys, interrupt, and ansi escape sequences. 
+- ⚡ **Interactive Command Handling**: Supports interactive commands using arrow keys, interrupt, and ansi escape sequences. 
+- ⚡ **REPL support**: [beta] Supports python/node and other REPL execution.
 
 ###  🪜 Steps: 
 1. Run the [cli client](https://github.com/rusiaaman/wcgw?tab=readme-ov-file#client) in any directory of choice.
@@ -103,7 +106,7 @@ Run the server
 If you don't have public ip and domain name, you can use `ngrok` or similar services to get a https address to the api.
 
 The specify the server url in the `wcgw` command like so
-`wcgw --server-url https://your-url/register`
+`wcgw --server-url https://your-url/v1/register`
 
 # [Optional] Local shell access with openai API key
 

wcgw-1.0.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+wcgw/__init__.py,sha256=VMi3gCAN_4_Ft8v5dY74u6bt5X-H1QhFJqTMZRd4fvk,76
+wcgw/types_.py,sha256=qpMRh1y136GkjIONIFowLy56v5p0OcdNqEH2mTmHPqU,756
+wcgw/client/__main__.py,sha256=ngI_vBcLAv7fJgmS4w4U7tuWtalGB8c7W5qebuT6Z6o,30
+wcgw/client/basic.py,sha256=2rY5pKm9dBBRuku0uIBUeLWjNc3iXjTSK9CmfuTxoS8,16196
+wcgw/client/claude.py,sha256=Bp45-UMBIJd-4tzX618nu-SpRbVtkTb1Es6c_gW6xy0,14861
+wcgw/client/common.py,sha256=grH-yV_4tnTQZ29xExn4YicGLxEq98z-HkEZwH0ReSg,1410
+wcgw/client/openai_adapters.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+wcgw/client/openai_utils.py,sha256=YNwCsA-Wqq7jWrxP0rfQmBTb1dI0s7dWXzQqyTzOZT4,2629
+wcgw/client/tools.py,sha256=UHPbpaQ70sFzE5_TYYyhV2FAE6BcbEgyRrTexaIpT6Y,21108
+wcgw/relay/serve.py,sha256=y2jGbZ0vKKONOD8TkJjSAaH8N6qGLa-rFtj5DxjUFIw,8916
+wcgw/relay/static/privacy.txt,sha256=s9qBdbx2SexCpC_z33sg16TptmAwDEehMCLz4L50JLc,529
+wcgw-1.0.0.dist-info/METADATA,sha256=8OScNwT6aecVULmjuk2oEPNIYGrhhMaAdybsdm57kYM,5217
+wcgw-1.0.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+wcgw-1.0.0.dist-info/entry_points.txt,sha256=WlIB825-Vm9ZtNzgENQsbHj4DRMkbpVR7uSkQyBlaPA,93
+wcgw-1.0.0.dist-info/RECORD,,

{wcgw-0.1.2.dist-info → wcgw-1.0.0.dist-info}/entry_points.txt

@@ -1,3 +1,4 @@
 [console_scripts]
 wcgw = wcgw:listen
 wcgw_local = wcgw:app
+wcgw_relay = wcgw.relay.serve:run

wcgw/__main__.py

@@ -1,3 +0,0 @@
-from wcgw.tools import run
-
-run()

wcgw-0.1.2.dist-info/RECORD

@@ -1,12 +0,0 @@
-wcgw/__init__.py,sha256=okSsOWpTKDjEQzgOin3Kdpx4Mc3MFX1RunjopHQSIWE,62
-wcgw/__main__.py,sha256=MjJnFwfYzA1rW47xuSP1EVsi53DTHeEGqESkQwsELFQ,34
-wcgw/basic.py,sha256=z1RVJMTDE1-J33nAPSfMZDdJBliSPCFh55SDCvtDLFI,16198
-wcgw/claude.py,sha256=Bp45-UMBIJd-4tzX618nu-SpRbVtkTb1Es6c_gW6xy0,14861
-wcgw/common.py,sha256=grH-yV_4tnTQZ29xExn4YicGLxEq98z-HkEZwH0ReSg,1410
-wcgw/openai_adapters.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-wcgw/openai_utils.py,sha256=YNwCsA-Wqq7jWrxP0rfQmBTb1dI0s7dWXzQqyTzOZT4,2629
-wcgw/tools.py,sha256=PDfSd6hcL60kaY947txztECWSDoGJrX_DBTz-bTfocY,16811
-wcgw-0.1.2.dist-info/METADATA,sha256=ramx2gvuo6NimIPSyN-vQPTAAXneklziWeBDNl9kXAA,5076
-wcgw-0.1.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-wcgw-0.1.2.dist-info/entry_points.txt,sha256=T-IH7w6Vc650hr8xksC8kJfbJR4uwN8HDudejwDwrNM,59
-wcgw-0.1.2.dist-info/RECORD,,