Adding reset shell api. Added possibility of having edit file tool

README.md

@@ -78,7 +78,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
 

gpt_action_json_schema.json

@@ -50,6 +50,46 @@
         }
       }
     },
+    "/v1/reset_shell": {
+      "post": {
+        "x-openai-isConsequential": false,
+        "summary": "Reset Shell",
+        "operationId": "reset_shell_v1_reset_shell_post",
+        "requestBody": {
+          "content": {
+            "application/json": {
+              "schema": {
+                "$ref": "#/components/schemas/ResetShellWithUUID"
+              }
+            }
+          },
+          "required": true
+        },
+        "responses": {
+          "200": {
+            "description": "Successful Response",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "type": "string",
+                  "title": "Response Reset Shell V1 Reset Shell Post"
+                }
+              }
+            }
+          },
+          "422": {
+            "description": "Validation Error",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/HTTPValidationError"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
     "/v1/bash_command": {
       "post": {
         "x-openai-isConsequential": false,
@@ -228,6 +268,29 @@
         "type": "object",
         "title": "HTTPValidationError"
       },
+      "ResetShellWithUUID": {
+        "properties": {
+          "should_reset": {
+            "type": "boolean",
+            "enum": [
+              true
+            ],
+            "const": true,
+            "title": "Should Reset",
+            "default": true
+          },
+          "user_id": {
+            "type": "string",
+            "format": "uuid",
+            "title": "User Id"
+          }
+        },
+        "type": "object",
+        "required": [
+          "user_id"
+        ],
+        "title": "ResetShellWithUUID"
+      },
       "ValidationError": {
         "properties": {
           "loc": {

gpt_instructions.txt

@@ -19,13 +19,17 @@ Instructions for `BashCommand`:
 
 Instructions for `Write File`
 - Write content to a file. Provide file path and content. Use this instead of BashCommand for writing files.
+- This doesn't create any directories, please create directories using `mkdir -p` BashCommand.
 - Important: all relative paths are relative to last CWD.
 
 Instructions for `BashInteraction`
 - Interact with running program using this tool
 - Special keys like arrows, interrupts, enter, etc.
 - Send text input to the running program.
 
+Instructions for `ResetShell`
+- Resets the shell. Use only if all interrupts and prompt reset attempts have failed repeatedly.
+
 ---
 
 Always critically think and debate with yourself to solve the problem. Understand the context and the code by reading as much resources as possible before writing a single piece of code.

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
     "websockets>=13.1",
     "pydantic>=2.9.2",
     "semantic-version>=2.10.0",
+    "nltk>=3.9.1",
 ]
 
 [project.urls]

src/wcgw/client/basic.py

@@ -20,7 +20,7 @@
 from typer import Typer
 import uuid
 
-from ..types_ import BashCommand, BashInteraction, ReadImage, Writefile
+from ..types_ import BashCommand, BashInteraction, ReadImage, Writefile, ResetShell
 
 from .common import Models, discard_input
 from .common import CostData, History
@@ -177,6 +177,10 @@ def loop(
         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

src/wcgw/client/tools.py

@@ -18,7 +18,6 @@
 )
 import uuid
 from pydantic import BaseModel, TypeAdapter
-import semantic_version
 import typer
 from websockets.sync.client import connect as syncconnect
 
@@ -41,8 +40,8 @@
     ChatCompletionMessage,
     ParsedChatCompletionMessage,
 )
-
-from ..types_ import Writefile
+from nltk.metrics.distance import edit_distance
+from ..types_ import FileEditFindReplace, ResetShell, Writefile
 
 from ..types_ import BashCommand
 
@@ -143,10 +142,21 @@ def _get_exit_code() -> int:
 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`"""
+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:
@@ -378,14 +388,9 @@ def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> T:
     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(PROMPT)
-        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")
@@ -411,6 +416,54 @@ 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")
+    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{out_string}\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"
+
+
 class DoneFlag(BaseModel):
     task_output: str
 
@@ -438,15 +491,19 @@ def which_tool(args: str) -> BaseModel:
         Confirmation
         | BashCommand
         | BashInteraction
+        | ResetShell
         | Writefile
+        | FileEditFindReplace
         | AIAssistant
         | DoneFlag
         | ReadImage
     ](
         Confirmation
         | BashCommand
         | BashInteraction
+        | ResetShell
         | Writefile
+        | FileEditFindReplace
         | AIAssistant
         | DoneFlag
         | ReadImage
@@ -459,7 +516,9 @@ def get_tool_output(
     | Confirmation
     | BashCommand
     | BashInteraction
+    | ResetShell
     | Writefile
+    | FileEditFindReplace
     | AIAssistant
     | DoneFlag
     | ReadImage,
@@ -473,15 +532,19 @@ def get_tool_output(
             Confirmation
             | BashCommand
             | BashInteraction
+            | ResetShell
             | Writefile
+            | FileEditFindReplace
             | AIAssistant
             | DoneFlag
             | ReadImage
         ](
             Confirmation
             | BashCommand
             | BashInteraction
+            | ResetShell
             | Writefile
+            | FileEditFindReplace
             | AIAssistant
             | DoneFlag
             | ReadImage
@@ -499,6 +562,9 @@ def get_tool_output(
     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
@@ -508,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}")
 
@@ -524,7 +593,7 @@ def get_tool_output(
 
 
 class Mdata(BaseModel):
-    data: BashCommand | BashInteraction | Writefile
+    data: BashCommand | BashInteraction | Writefile | ResetShell | FileEditFindReplace
 
 
 execution_lock = threading.Lock()