Merge pull request #3 from rusiaaman/feat/repl

Feat/repl

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.
+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.

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 authors = [{ name = "Aman Rusia", email = "[email protected]" }]
 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"

src/wcgw/basic.py

@@ -159,14 +159,15 @@ 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 <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.
 - 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(

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}'"