diff --git a/gpt_instructions.txt b/gpt_instructions.txt index 73d27c8..3f538ea 100644 --- a/gpt_instructions.txt +++ b/gpt_instructions.txt @@ -19,6 +19,7 @@ Instructions for `Execute Bash`: - 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. +- You can run python/node/other REPL code lines using `execute_command` too. NOTE: `execute_command` doesn't create a new shell, it uses the same shell. Instructions for `Write File` - Write content to a file. Provide file path and content. Use this instead of ExecuteBash for writing files. @@ -28,4 +29,10 @@ Instructions for `Write File` 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. --- -Ask the user for the user_id `UUID` if they haven't provided in the first message. \ No newline at end of file +Ask the user for the user_id `UUID` if they haven't provided in the first message. + +--- +Error references: +1. "Input should be a valid integer" + You are probably send_ascii command a string. If you are in REPL mode, use `execute_command` instead. + Otherwise convert the string into sequence of ascii integers. \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 102c345..0a985c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] authors = [{ name = "Aman Rusia", email = "gapypi@arcfu.com" }] name = "wcgw" -version = "0.1.2" +version = "0.2.0" description = "What could go wrong giving full shell access to chatgpt?" readme = "README.md" requires-python = ">=3.10, <3.13" diff --git a/src/wcgw/basic.py b/src/wcgw/basic.py index 9fc36e2..ca27b5f 100644 --- a/src/wcgw/basic.py +++ b/src/wcgw/basic.py @@ -159,7 +159,7 @@ def loop( ExecuteBash, description=""" - Execute a bash script. This is stateful (beware with subsequent calls). -- Execute commands using `execute_command` attribute. +- Execute commands using `execute_command` attribute. You can run python/node/other REPL code lines using `execute_command` too. - Do not use interactive commands like nano. Prefer writing simpler commands. - Last line will always be `(exit )` 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]`. @@ -167,6 +167,7 @@ def loop( - 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. +- You can run python/node/other REPL code lines using `execute_command` too. NOTE: `execute_command` doesn't create a new shell, it uses the same shell. """, ), openai.pydantic_function_tool( diff --git a/src/wcgw/tools.py b/src/wcgw/tools.py index 7ddbbaf..f249802 100644 --- a/src/wcgw/tools.py +++ b/src/wcgw/tools.py @@ -2,6 +2,7 @@ import base64 import json import mimetypes +import re import sys import threading import traceback @@ -77,17 +78,20 @@ class Writefile(BaseModel): file_content: str +PROMPT = "#@@" + + def start_shell() -> pexpect.spawn: 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 +107,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() @@ -136,17 +146,51 @@ class ExecuteBash(BaseModel): 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]` +1. Get its output using `send_ascii: [10] or send_ascii: ["Enter"]` 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`""" +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 + + def execute_bash( enc: tiktoken.Encoding, bash_arg: ExecuteBash, max_tokens: Optional[int] ) -> tuple[str, float]: global SHELL, BASH_STATE try: + is_interrupt = False if bash_arg.execute_command: + updated_repl_mode = update_repl_prompt(bash_arg.execute_command) + if updated_repl_mode: + BASH_STATE = "running" + response = "Prompt updated, you can execute REPL lines using execute_command now" + console.print(response) + return ( + response, + 0, + ) + + console.print(f"$ {bash_arg.execute_command}") if BASH_STATE == "waiting_for_input": raise ValueError(WAITING_INPUT_MESSAGE) elif BASH_STATE == "wont_exit": @@ -160,14 +204,14 @@ def execute_bash( 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)) + if char == 3: + is_interrupt = True if char == "Key-up": SHELL.send("\033[A") elif char == "Key-down": @@ -180,6 +224,7 @@ def execute_bash( SHELL.send("\n") elif char == "Ctrl-c": SHELL.sendintr() + is_interrupt = True else: raise Exception("Nothing to send") BASH_STATE = "running" @@ -190,16 +235,11 @@ def execute_bash( raise wait = 5 - index = SHELL.expect(["#@@", pexpect.TIMEOUT], timeout=wait) - running = "" - while index == 1: - if wait > TIMEOUT: - raise TimeoutError("Timeout while waiting for shell prompt") - + index = SHELL.expect([PROMPT, pexpect.TIMEOUT], timeout=wait) + if index == 1: BASH_STATE = "waiting_for_input" text = SHELL.before or "" - print(text[len(running) :]) - running = text + print(text) text = render_terminal_output(text) tokens = enc.encode(text) @@ -208,7 +248,21 @@ def execute_bash( text = "...(truncated)\n" + enc.decode(tokens[-(max_tokens - 1) :]) last_line = "(pending)" - return text + f"\n{last_line}", 0 + text = text + f"\n{last_line}" + + if is_interrupt: + text = ( + text + + """ +Failure interrupting. Have you entered a new REPL like python, node, ipython, etc.? Or have you exited from a previous REPL program? +If yes: + Run execute_command: "wcgw_update_prompt()" to enter the new REPL mode. +If no: + Try Ctrl-c or Ctrl-d again. +""" + ) + + return text, 0 assert isinstance(SHELL.before, str) output = render_terminal_output(SHELL.before) @@ -284,7 +338,7 @@ def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> T: def read_image_from_shell(file_path: str) -> ImageData: if not os.path.isabs(file_path): SHELL.sendline("pwd") - SHELL.expect("#@@") + 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) @@ -303,7 +357,7 @@ def read_image_from_shell(file_path: str) -> ImageData: def write_file(writefile: Writefile) -> str: if not os.path.isabs(writefile.file_path): SHELL.sendline("pwd") - SHELL.expect("#@@") + SHELL.expect(PROMPT) 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}'" diff --git a/uv.lock b/uv.lock index 14005f0..2478a97 100644 --- a/uv.lock +++ b/uv.lock @@ -876,7 +876,7 @@ wheels = [ [[package]] name = "wcgw" -version = "0.0.10" +version = "0.1.2" source = { editable = "." } dependencies = [ { name = "fastapi" },