wcgw 0.1.0__py3-none-any.whl

wcgw/__init__.py

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

wcgw/__main__.py

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

wcgw/basic.py

@@ -0,0 +1,341 @@
+import json
+from pathlib import Path
+import sys
+import traceback
+from typing import Callable, DefaultDict, Optional, cast
+import openai
+from openai import OpenAI
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessage,
+    ParsedChatCompletionMessage,
+)
+import rich
+from typer import Typer
+import uuid
+
+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, GetShellOutputLastCommand
+
+from .tools import (
+    BASH_CLF_OUTPUT,
+    Confirmation,
+    DoneFlag,
+    Writefile,
+    get_is_waiting_user_input,
+    get_tool_output,
+    SHELL,
+    start_shell,
+    which_tool,
+)
+import tiktoken
+
+from urllib import parse
+import subprocess
+import os
+import tempfile
+
+import toml
+from pydantic import BaseModel
+
+
+from dotenv import load_dotenv
+
+
+class Config(BaseModel):
+    model: Models
+    secondary_model: Models
+    cost_limit: float
+    cost_file: dict[Models, CostData]
+    cost_unit: str = "$"
+
+
+def text_from_editor(console: rich.console.Console) -> str:
+    # First consume all the input till now
+    discard_input()
+    console.print("\n---------------------------------------\n# User message")
+    data = input()
+    if data:
+        return data
+    editor = os.environ.get("EDITOR", "vim")
+    with tempfile.NamedTemporaryFile(suffix=".tmp") as tf:
+        subprocess.run([editor, tf.name], check=True)
+        with open(tf.name, "r") as f:
+            data = f.read()
+            console.print(data)
+            return data
+
+
+def save_history(history: History, session_id: str) -> None:
+    myid = str(history[1]["content"]).replace("/", "_").replace(" ", "_").lower()[:60]
+    myid += "_" + session_id
+    myid = myid + ".json"
+
+    mypath = Path(".wcgw") / myid
+    mypath.parent.mkdir(parents=True, exist_ok=True)
+    with open(mypath, "w") as f:
+        json.dump(history, f, indent=3)
+
+
+app = Typer(pretty_exceptions_show_locals=False)
+
+
+@app.command()
+def loop(
+    first_message: Optional[str] = None,
+    limit: Optional[float] = None,
+    resume: Optional[str] = None,
+) -> tuple[str, float]:
+    load_dotenv()
+
+    session_id = str(uuid.uuid4())[:6]
+
+    history: History = []
+    if resume:
+        if resume == "latest":
+            resume_path = sorted(Path(".wcgw").iterdir(), key=os.path.getmtime)[-1]
+        else:
+            resume_path = Path(resume)
+        if not resume_path.exists():
+            raise FileNotFoundError(f"File {resume} not found")
+        with resume_path.open() as f:
+            history = json.load(f)
+        if len(history) <= 2:
+            raise ValueError("Invalid history file")
+        if history[1]["role"] != "user":
+            raise ValueError("Invalid history file, second message should be user")
+        first_message = ""
+
+    my_dir = os.path.dirname(__file__)
+    config_file = os.path.join(my_dir, "..", "..", "config.toml")
+    with open(config_file) as f:
+        config_json = toml.load(f)
+        config = Config.model_validate(config_json)
+
+    if limit is not None:
+        config.cost_limit = limit
+    limit = config.cost_limit
+
+    enc = tiktoken.encoding_for_model(
+        config.model if not config.model.startswith("o1") else "gpt-4o"
+    )
+    is_waiting_user_input = get_is_waiting_user_input(
+        config.model, config.cost_file[config.model]
+    )
+
+    tools = [
+        openai.pydantic_function_tool(
+            ExecuteBash,
+            description="""
+Execute a bash script. Stateful (beware with subsequent calls).
+Execute commands using `execute_command` attribute.
+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 `(waiting for input)` which will be the case if you've run any interactive command (which you shouldn't run) by mistake. You can then send input using `send_ascii` attributes.
+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.""",
+        ),
+        openai.pydantic_function_tool(
+            GetShellOutputLastCommand,
+            description="Get output of the last command run in the shell. Use this in case you want to know status of a running program.",
+        ),
+        openai.pydantic_function_tool(
+            Writefile,
+            description="Write content to a file. Provide file path and content. Use this instead of ExecuteBash for writing files.",
+        ),
+    ]
+    uname_sysname = os.uname().sysname
+    uname_machine = os.uname().machine
+
+    system = f"""
+You're a cli assistant.
+
+Instructions:
+
+    - You should use the provided bash execution tool to run script to complete objective. 
+    - Do not use sudo. Do not use interactive commands. 
+    - Ask user for confirmation before running anything major
+    
+System information:
+    - System: {uname_sysname}
+    - Machine: {uname_machine}
+"""
+
+    has_tool_output = False
+    if not history:
+        history = [{"role": "system", "content": system}]
+    else:
+        if history[-1]["role"] == "tool":
+            has_tool_output = True
+
+    client = OpenAI()
+
+    cost: float = 0
+    input_toks = 0
+    output_toks = 0
+    system_console = rich.console.Console(style="blue", highlight=False)
+    error_console = rich.console.Console(style="red", highlight=False)
+    user_console = rich.console.Console(style="bright_black", highlight=False)
+    assistant_console = rich.console.Console(style="white bold", highlight=False)
+
+    while True:
+        if cost > limit:
+            system_console.print(
+                f"\nCost limit exceeded. Current cost: {cost}, input tokens: {input_toks}, output tokens: {output_toks}"
+            )
+            break
+
+        if not has_tool_output:
+            if first_message:
+                msg = first_message
+                first_message = ""
+            else:
+                msg = text_from_editor(user_console)
+
+            history.append({"role": "user", "content": msg})
+        else:
+            has_tool_output = False
+
+        cost_, input_toks_ = get_input_cost(
+            config.cost_file[config.model], enc, history
+        )
+        cost += cost_
+        input_toks += input_toks_
+
+        stream = client.chat.completions.create(
+            messages=history,
+            model=config.model,
+            stream=True,
+            tools=tools,
+        )
+
+        system_console.print(
+            "\n---------------------------------------\n# Assistant response",
+            style="bold",
+        )
+        tool_call_args_by_id = DefaultDict[str, DefaultDict[int, str]](
+            lambda: DefaultDict(str)
+        )
+        _histories: History = []
+        item: ChatCompletionMessageParam
+        full_response: str = ""
+        try:
+            for chunk in stream:
+                if chunk.choices[0].finish_reason == "tool_calls":
+                    assert tool_call_args_by_id
+                    item = {
+                        "role": "assistant",
+                        "content": full_response,
+                        "tool_calls": [
+                            {
+                                "id": tool_call_id + str(toolindex),
+                                "type": "function",
+                                "function": {
+                                    "arguments": tool_args,
+                                    "name": "execute_bash",
+                                },
+                            }
+                            for tool_call_id, toolcallargs in tool_call_args_by_id.items()
+                            for toolindex, tool_args in toolcallargs.items()
+                        ],
+                    }
+                    cost_, output_toks_ = get_output_cost(
+                        config.cost_file[config.model], enc, item
+                    )
+                    cost += cost_
+                    system_console.print(
+                        f"\n---------------------------------------\n# Assistant invoked tools: {[which_tool(tool['function']['arguments']) for tool in item['tool_calls']]}"
+                    )
+                    system_console.print(f"\nTotal cost: {config.cost_unit}{cost:.3f}")
+                    output_toks += output_toks_
+
+                    _histories.append(item)
+                    for tool_call_id, toolcallargs in tool_call_args_by_id.items():
+                        for toolindex, tool_args in toolcallargs.items():
+                            try:
+                                output_or_done, cost_ = get_tool_output(
+                                    json.loads(tool_args),
+                                    enc,
+                                    limit - cost,
+                                    loop,
+                                    is_waiting_user_input,
+                                )
+                            except Exception as e:
+                                output_or_done = (
+                                    f"GOT EXCEPTION while calling tool. Error: {e}"
+                                )
+                                tb = traceback.format_exc()
+                                error_console.print(output_or_done + "\n" + tb)
+                                cost_ = 0
+                            cost += cost_
+                            system_console.print(
+                                f"\nTotal cost: {config.cost_unit}{cost:.3f}"
+                            )
+
+                            if isinstance(output_or_done, DoneFlag):
+                                system_console.print(
+                                    f"\n# Task marked done, with output {output_or_done.task_output}",
+                                )
+                                system_console.print(
+                                    f"\nTotal cost: {config.cost_unit}{cost:.3f}"
+                                )
+                                return output_or_done.task_output, cost
+                            output = output_or_done
+
+                            item = {
+                                "role": "tool",
+                                "content": str(output),
+                                "tool_call_id": tool_call_id + str(toolindex),
+                            }
+                            cost_, output_toks_ = get_output_cost(
+                                config.cost_file[config.model], enc, item
+                            )
+                            cost += cost_
+                            output_toks += output_toks_
+
+                            _histories.append(item)
+                    has_tool_output = True
+                    break
+                elif chunk.choices[0].finish_reason:
+                    assistant_console.print("")
+                    item = {
+                        "role": "assistant",
+                        "content": full_response,
+                    }
+                    cost_, output_toks_ = get_output_cost(
+                        config.cost_file[config.model], enc, item
+                    )
+                    cost += cost_
+                    output_toks += output_toks_
+
+                    system_console.print(f"\nTotal cost: {config.cost_unit}{cost:.3f}")
+                    _histories.append(item)
+                    break
+
+                if chunk.choices[0].delta.tool_calls:
+                    tool_call = chunk.choices[0].delta.tool_calls[0]
+                    if tool_call.function and tool_call.function.arguments:
+                        tool_call_args_by_id[tool_call.id or ""][tool_call.index] += (
+                            tool_call.function.arguments
+                        )
+
+                chunk_str = chunk.choices[0].delta.content or ""
+                assistant_console.print(chunk_str, end="")
+                full_response += chunk_str
+        except KeyboardInterrupt:
+            has_tool_output = False
+            input("Interrupted...enter to redo the current turn")
+        else:
+            history.extend(_histories)
+
+            save_history(history, session_id)
+
+    return "Couldn't finish the task", cost
+
+
+if __name__ == "__main__":
+    app()

wcgw/common.py

@@ -0,0 +1,47 @@
+import select
+import sys
+import termios
+import tty
+from typing import Literal
+from pydantic import BaseModel
+
+
+class CostData(BaseModel):
+    cost_per_1m_input_tokens: float
+    cost_per_1m_output_tokens: float
+
+
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessage,
+    ParsedChatCompletionMessage,
+)
+
+History = list[ChatCompletionMessageParam]
+Models = Literal["gpt-4o-2024-08-06", "gpt-4o-mini"]
+
+
+def discard_input() -> None:
+    # Get the file descriptor for stdin
+    fd = sys.stdin.fileno()
+
+    # Save current terminal settings
+    old_settings = termios.tcgetattr(fd)
+
+    try:
+        # Switch terminal to non-canonical mode where input is read immediately
+        tty.setcbreak(fd)
+
+        # Discard all input
+        while True:
+            # Check if there is input to be read
+            if sys.stdin in select.select([sys.stdin], [], [], 0)[0]:
+                sys.stdin.read(
+                    1
+                )  # Read one character at a time to flush the input buffer
+            else:
+                break
+    finally:
+        # Restore old terminal settings
+        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)

wcgw/openai_utils.py

@@ -0,0 +1,68 @@
+import json
+from pathlib import Path
+import select
+import sys
+import termios
+import traceback
+import tty
+from typing import Callable, DefaultDict, Literal, Optional, cast
+import openai
+from openai import OpenAI
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessage,
+    ParsedChatCompletionMessage,
+)
+import rich
+import tiktoken
+from typer import Typer
+import uuid
+
+from .common import CostData, History
+
+
+def get_input_cost(
+    cost_map: CostData, enc: tiktoken.Encoding, history: History
+) -> tuple[float, int]:
+    input_tokens = 0
+    for msg in history:
+        content = msg["content"]
+        if not isinstance(content, str):
+            raise ValueError(f"Expected content to be string, got {type(content)}")
+        input_tokens += len(enc.encode(content))
+    cost = input_tokens * cost_map.cost_per_1m_input_tokens / 1_000_000
+    return cost, input_tokens
+
+
+def get_output_cost(
+    cost_map: CostData,
+    enc: tiktoken.Encoding,
+    item: ChatCompletionMessage | ChatCompletionMessageParam,
+) -> tuple[float, int]:
+    if isinstance(item, ChatCompletionMessage):
+        content = item.content
+        if not isinstance(content, str):
+            raise ValueError(f"Expected content to be string, got {type(content)}")
+    else:
+        if not isinstance(item["content"], str):
+            raise ValueError(
+                f"Expected content to be string, got {type(item['content'])}"
+            )
+        content = item["content"]
+        if item["role"] == "tool":
+            return 0, 0
+    output_tokens = len(enc.encode(content))
+
+    if "tool_calls" in item:
+        item = cast(ChatCompletionAssistantMessageParam, item)
+        toolcalls = item["tool_calls"]
+        for tool_call in toolcalls or []:
+            output_tokens += len(enc.encode(tool_call["function"]["arguments"]))
+    elif isinstance(item, ParsedChatCompletionMessage):
+        if item.tool_calls:
+            for tool_callf in item.tool_calls:
+                output_tokens += len(enc.encode(tool_callf.function.arguments))
+
+    cost = output_tokens * cost_map.cost_per_1m_output_tokens / 1_000_000
+    return cost, output_tokens

wcgw/tools.py

@@ -0,0 +1,465 @@
+import asyncio
+import json
+import sys
+import threading
+import traceback
+from typing import Callable, Literal, Optional, ParamSpec, Sequence, TypeVar, TypedDict
+import uuid
+from pydantic import BaseModel, TypeAdapter
+
+import os
+import tiktoken
+import petname  # type: ignore[import]
+import pexpect
+from typer import Typer
+import websockets
+
+import rich
+import pyte
+from dotenv import load_dotenv
+
+import openai
+from openai import OpenAI
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessage,
+    ParsedChatCompletionMessage,
+)
+
+from .common import CostData, Models, discard_input
+
+from .openai_utils import get_input_cost, get_output_cost
+
+console = rich.console.Console(style="magenta", highlight=False)
+
+TIMEOUT = 30
+
+
+def render_terminal_output(text: str) -> str:
+    screen = pyte.Screen(160, 500)
+    screen.set_mode(pyte.modes.LNM)
+    stream = pyte.Stream(screen)
+    stream.feed(text)
+    # Filter out empty lines
+    dsp = screen.display[::-1]
+    for i, line in enumerate(dsp):
+        if line.strip():
+            break
+    else:
+        i = len(dsp)
+    return "\n".join(screen.display[: len(dsp) - i])
+
+
+class Confirmation(BaseModel):
+    prompt: str
+
+
+def ask_confirmation(prompt: Confirmation) -> str:
+    response = input(prompt.prompt + " [y/n] ")
+    return "Yes" if response.lower() == "y" else "No"
+
+
+class Writefile(BaseModel):
+    file_path: str
+    file_content: str
+
+
+def start_shell():
+    SHELL = pexpect.spawn(
+        "/bin/bash",
+        env={**os.environ, **{"PS1": "#@@"}},
+        echo=False,
+        encoding="utf-8",
+        timeout=TIMEOUT,
+    )  # type: ignore[arg-type]
+    SHELL.expect("#@@")
+    SHELL.sendline("stty -icanon -echo")
+    SHELL.expect("#@@")
+    return SHELL
+
+
+SHELL = start_shell()
+
+
+def _get_exit_code() -> int:
+    SHELL.sendline("echo $?")
+    SHELL.expect("#@@")
+    assert isinstance(SHELL.before, str)
+    return int((SHELL.before))
+
+
+Specials = Literal["Key-up", "Key-down", "Key-left", "Key-right", "Enter", "Ctrl-c"]
+
+
+class ExecuteBash(BaseModel):
+    execute_command: Optional[str] = None
+    send_ascii: Optional[Sequence[int | Specials]] = None
+
+
+class GetShellOutputLastCommand(BaseModel):
+    type: Literal["get_output_of_last_command"] = "get_output_of_last_command"
+
+
+BASH_CLF_OUTPUT = Literal["running", "waiting_for_input", "wont_exit"]
+BASH_STATE: BASH_CLF_OUTPUT = "running"
+
+
+def get_output_of_last_command(enc: tiktoken.Encoding) -> str:
+    global SHELL, BASH_STATE
+    output = render_terminal_output(SHELL.before)
+
+    tokens = enc.encode(output)
+    if len(tokens) >= 2048:
+        output = "...(truncated)\n" + enc.decode(tokens[-2047:])
+
+    return output
+
+
+WETTING_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 `GetShellOutputLastCommand` OR
+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 execute_bash(
+    enc: tiktoken.Encoding,
+    bash_arg: ExecuteBash,
+    is_waiting_user_input: Callable[[str], tuple[BASH_CLF_OUTPUT, float]],
+) -> tuple[str, float]:
+    global SHELL, BASH_STATE
+    try:
+        if bash_arg.execute_command:
+            if BASH_STATE == "waiting_for_input":
+                raise ValueError(WETTING_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`"""
+                )
+            command = bash_arg.execute_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))
+                if char == "Key-up":
+                    SHELL.send("\033[A")
+                elif char == "Key-down":
+                    SHELL.send("\033[B")
+                elif char == "Key-left":
+                    SHELL.send("\033[D")
+                elif char == "Key-right":
+                    SHELL.send("\033[C")
+                elif char == "Enter":
+                    SHELL.send("\n")
+                elif char == "Ctrl-c":
+                    SHELL.sendintr()
+        else:
+            raise Exception("Nothing to send")
+        BASH_STATE = "running"
+
+    except KeyboardInterrupt:
+        SHELL.close(True)
+        SHELL = start_shell()
+        raise
+
+    wait = timeout = 5
+    index = SHELL.expect(["#@@", pexpect.TIMEOUT], timeout=wait)
+    running = ""
+    while index == 1:
+        if wait > TIMEOUT:
+            raise TimeoutError("Timeout while waiting for shell prompt")
+
+        text = SHELL.before
+        print(text[len(running) :])
+        running = text
+
+        text = render_terminal_output(text)
+        BASH_STATE, cost = is_waiting_user_input(text)
+        if BASH_STATE == "waiting_for_input" or BASH_STATE == "wont_exit":
+            tokens = enc.encode(text)
+
+            if len(tokens) >= 2048:
+                text = "...(truncated)\n" + enc.decode(tokens[-2047:])
+
+            last_line = (
+                "(waiting for input)"
+                if BASH_STATE == "waiting_for_input"
+                else "(won't exit)"
+            )
+            return text + f"\n{last_line}", cost
+        index = SHELL.expect(["#@@", pexpect.TIMEOUT], timeout=wait)
+        wait += timeout
+
+    assert isinstance(SHELL.before, str)
+    output = render_terminal_output(SHELL.before)
+
+    tokens = enc.encode(output)
+    if len(tokens) >= 2048:
+        output = "...(truncated)\n" + enc.decode(tokens[-2047:])
+
+    try:
+        exit_code = _get_exit_code()
+        output += f"\n(exit {exit_code})"
+
+    except ValueError:
+        console.print("Malformed output, restarting shell", style="red")
+        # Malformed output, restart shell
+        SHELL.close(True)
+        SHELL = start_shell()
+        output = "(exit shell has restarted)"
+    return output, 0
+
+
+Param = ParamSpec("Param")
+
+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":
+            raise ValueError(WETTING_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 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 = SHELL.before.strip()
+        writefile.file_path = os.path.join(current_dir, writefile.file_path)
+    os.makedirs(os.path.dirname(writefile.file_path), exist_ok=True)
+    try:
+        with open(writefile.file_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}")
+    return "Success"
+
+
+class DoneFlag(BaseModel):
+    task_output: str
+
+
+def mark_finish(done: DoneFlag) -> DoneFlag:
+    return done
+
+
+class AIAssistant(BaseModel):
+    instruction: str
+    desired_output: str
+
+
+def take_help_of_ai_assistant(
+    aiassistant: AIAssistant,
+    limit: float,
+    loop_call: Callable[[str, float], tuple[str, float]],
+) -> tuple[str, float]:
+    output, cost = loop_call(aiassistant.instruction, limit)
+    return output, cost
+
+
+class AddTasks(BaseModel):
+    task_statement: str
+
+
+def add_task(addtask: AddTasks) -> str:
+    petname_id = petname.Generate(2, "-")
+    return petname_id
+
+
+class RemoveTask(BaseModel):
+    task_id: str
+
+
+def remove_task(removetask: RemoveTask) -> str:
+    return "removed"
+
+
+def which_tool(args: str) -> BaseModel:
+    adapter = TypeAdapter[
+        Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag
+    ](Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag)
+    return adapter.validate_python(json.loads(args))
+
+
+def get_tool_output(
+    args: dict | Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag,
+    enc: tiktoken.Encoding,
+    limit: float,
+    loop_call: Callable[[str, float], tuple[str, float]],
+    is_waiting_user_input: Callable[[str], tuple[BASH_CLF_OUTPUT, float]],
+) -> tuple[str | DoneFlag, float]:
+    if isinstance(args, dict):
+        adapter = TypeAdapter[
+            Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag
+        ](Confirmation | ExecuteBash | Writefile | AIAssistant | DoneFlag)
+        arg = adapter.validate_python(args)
+    else:
+        arg = args
+    output: tuple[str | DoneFlag, float]
+    if isinstance(arg, Confirmation):
+        console.print("Calling ask confirmation tool")
+        output = ask_confirmation(arg), 0.0
+    elif isinstance(arg, ExecuteBash):
+        console.print("Calling execute bash tool")
+        output = execute_bash(enc, arg, is_waiting_user_input)
+    elif isinstance(arg, Writefile):
+        console.print("Calling write file tool")
+        output = write_file(arg), 0
+    elif isinstance(arg, DoneFlag):
+        console.print("Calling mark finish tool")
+        output = mark_finish(arg), 0.0
+    elif isinstance(arg, AIAssistant):
+        console.print("Calling AI assistant tool")
+        output = take_help_of_ai_assistant(arg, limit, loop_call)
+    elif isinstance(arg, AddTasks):
+        console.print("Calling add task tool")
+        output = add_task(arg), 0
+    elif isinstance(arg, get_output_of_last_command):
+        console.print("Calling get output of last program tool")
+        output = get_output_of_last_command(enc), 0
+    else:
+        raise ValueError(f"Unknown tool: {arg}")
+
+    console.print(str(output[0]))
+    return output
+
+
+History = list[ChatCompletionMessageParam]
+
+
+def get_is_waiting_user_input(model: Models, cost_data: CostData):
+    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 false 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 `normal` 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)
+curr_cost = 0.0
+
+
+class Mdata(BaseModel):
+    data: ExecuteBash | Writefile
+
+
+execution_lock = threading.Lock()
+
+
+def execute_user_input() -> None:
+    while True:
+        discard_input()
+        user_input = input()
+        if user_input:
+            with execution_lock:
+                try:
+                    console.log(
+                        execute_bash(
+                            default_enc,
+                            ExecuteBash(
+                                send_ascii=[ord(x) for x in user_input] + [ord("\n")]
+                            ),
+                            lambda x: ("wont_exit", 0),
+                        )[0]
+                    )
+                except Exception as e:
+                    traceback.print_exc()
+                    console.log(f"Error: {e}")
+
+
+async def register_client(server_url: str) -> None:
+    global default_enc, default_model, curr_cost
+    # Generate a unique UUID for this client
+    client_uuid = str(uuid.uuid4())
+    print(f"Connecting with UUID: {client_uuid}")
+
+    # Create the WebSocket connection
+    async with websockets.connect(f"{server_url}/{client_uuid}") as websocket:
+        try:
+            while True:
+                # Wait to receive data from the server
+                message = await websocket.recv()
+                print(message, type(message))
+                mdata = Mdata.model_validate_json(message)
+                with execution_lock:
+                    is_waiting_user_input = get_is_waiting_user_input(
+                        default_model, default_cost
+                    )
+                    try:
+                        output, cost = get_tool_output(
+                            mdata.data,
+                            default_enc,
+                            0.0,
+                            lambda x, y: ("", 0),
+                            is_waiting_user_input,
+                        )
+                        curr_cost += cost
+                        print(f"{curr_cost=}")
+                    except Exception as e:
+                        output = f"GOT EXCEPTION while calling tool. Error: {e}"
+                        traceback.print_exc()
+                    assert not isinstance(output, DoneFlag)
+                    await websocket.send(output)
+
+        except websockets.ConnectionClosed:
+            print(f"Connection closed for UUID: {client_uuid}")
+
+
+def run() -> None:
+    if len(sys.argv) > 1:
+        server_url = sys.argv[1]
+    else:
+        server_url = "ws://localhost:8000/register"
+    asyncio.run(register_client(server_url))

wcgw-0.1.0.dist-info/METADATA

@@ -0,0 +1,23 @@
+Metadata-Version: 2.3
+Name: wcgw
+Version: 0.1.0
+Summary: What could go wrong giving full shell access to chatgpt?
+Requires-Python: >=3.12
+Requires-Dist: fastapi>=0.115.0
+Requires-Dist: mypy>=1.11.2
+Requires-Dist: openai>=1.46.0
+Requires-Dist: petname>=2.6
+Requires-Dist: pexpect>=4.9.0
+Requires-Dist: pyte>=0.8.2
+Requires-Dist: python-dotenv>=1.0.1
+Requires-Dist: rich>=13.8.1
+Requires-Dist: shell>=1.0.1
+Requires-Dist: tiktoken==0.7.0
+Requires-Dist: toml>=0.10.2
+Requires-Dist: typer>=0.12.5
+Requires-Dist: types-pexpect>=4.9.0.20240806
+Requires-Dist: uvicorn>=0.31.0
+Requires-Dist: websockets>=13.1
+Description-Content-Type: text/markdown
+
+# What could go wrong giving full shell access to Chatgpt?

wcgw-0.1.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+wcgw/__init__.py,sha256=okSsOWpTKDjEQzgOin3Kdpx4Mc3MFX1RunjopHQSIWE,62
+wcgw/__main__.py,sha256=MjJnFwfYzA1rW47xuSP1EVsi53DTHeEGqESkQwsELFQ,34
+wcgw/basic.py,sha256=BiVjIwrtiz93SkUedDXSwtfVMKoV8-zEWeFKBIamVSQ,12372
+wcgw/common.py,sha256=jn39zTpaFUO1PWof_z7qBmklaZH5G1blzjlBvez0cg4,1225
+wcgw/openai_adapters.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+wcgw/openai_utils.py,sha256=4Hr9S2-WT8xhDdu3b2YVoX1l9AVxwCFdi_GJbQAx7Us,2202
+wcgw/tools.py,sha256=k6Xoaq_kqLnC6cgBQ9NGovwNDtam5zV7DPTNmq9zSbo,15069
+wcgw-0.1.0.dist-info/METADATA,sha256=oM2m4AYiPXrCPytISxbWfGWCMcS3t7lgw0TiitTIb-s,701
+wcgw-0.1.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+wcgw-0.1.0.dist-info/entry_points.txt,sha256=T-IH7w6Vc650hr8xksC8kJfbJR4uwN8HDudejwDwrNM,59
+wcgw-0.1.0.dist-info/RECORD,,

wcgw-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.25.0
+Root-Is-Purelib: true
+Tag: py3-none-any

wcgw-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+wcgw = wcgw:listen
+wcgw_local = wcgw:app