wcgw 2.8.10__py3-none-any.whl → 3.0.1__py3-none-any.whl

wcgw_cli/cli.py

@@ -16,7 +16,6 @@ def loop(
     first_message: Optional[str] = None,
     limit: Optional[float] = None,
     resume: Optional[str] = None,
-    computer_use: bool = False,
     version: bool = typer.Option(False, "--version", "-v"),
 ) -> tuple[str, float]:
     if version:
@@ -28,7 +27,6 @@ def loop(
             first_message=first_message,
             limit=limit,
             resume=resume,
-            computer_use=computer_use,
         )
     else:
         return openai_loop(

wcgw_cli/openai_client.py

@@ -23,25 +23,18 @@ from openai.types.chat import (
 from pydantic import BaseModel
 from typer import Typer
 
+from wcgw.client.bash_state.bash_state import BashState
 from wcgw.client.common import CostData, History, Models, discard_input
 from wcgw.client.memory import load_memory
+from wcgw.client.tool_prompts import TOOL_PROMPTS
 from wcgw.client.tools import (
-    DoneFlag,
+    Context,
     ImageData,
     default_enc,
     get_tool_output,
     initialize,
     which_tool,
-)
-from wcgw.types_ import (
-    BashCommand,
-    BashInteraction,
-    ContextSave,
-    FileEdit,
-    ReadFiles,
-    ReadImage,
-    ResetShell,
-    WriteIfEmpty,
+    which_tool_name,
 )
 
 from .openai_utils import get_input_cost, get_output_cost
@@ -114,7 +107,6 @@ def loop(
     first_message: Optional[str] = None,
     limit: Optional[float] = None,
     resume: Optional[str] = None,
-    computer_use: bool = False,
 ) -> tuple[str, float]:
     load_dotenv()
 
@@ -129,8 +121,8 @@ def loop(
             _, memory, _ = load_memory(
                 resume,
                 8000,
-                lambda x: default_enc.encode(x).ids,
-                lambda x: default_enc.decode(x),
+                lambda x: default_enc.encoder(x),
+                lambda x: default_enc.decoder(x),
             )
         except OSError:
             if resume == "latest":
@@ -167,89 +159,11 @@ def loop(
 
     tools = [
         openai.pydantic_function_tool(
-            BashCommand,
-            description="""
-- Execute a bash command. This is stateful (beware with subsequent calls).
-- Do not use interactive commands like nano. Prefer writing simpler commands.
-- 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 or the last 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.
-- The control will return to you in 5 seconds regardless of the status. For heavy commands, keep checking status using BashInteraction till they are finished.
-- Run long running commands in background using screen instead of "&".
-- Do not use 'cat' to read files, use ReadFiles tool instead.
-""",
-        ),
-        openai.pydantic_function_tool(
-            BashInteraction,
-            description="""
-- Interact with running program using this tool
-- Special keys like arrows, interrupts, enter, etc.
-- Send text input to the running program.
-- Send send_specials=["Enter"] to recheck status of a running program.
-- Only one of send_text, send_specials, send_ascii should be provided.""",
-        ),
-        openai.pydantic_function_tool(
-            ReadFiles,
-            description="""
-- Read full file content of one or more files.
-- Provide absolute file paths only
-""",
-        ),
-        openai.pydantic_function_tool(
-            WriteIfEmpty,
-            description="""
-- Write content to an empty or non-existent file. Provide file path and content. Use this instead of BashCommand for writing new files.
-- Provide absolute file path only.
-- For editing existing files, use FileEdit instead of this tool.""",
-        ),
-        openai.pydantic_function_tool(
-            FileEdit,
-            description="""
-- Use absolute file path only.
-- Use ONLY SEARCH/REPLACE blocks to edit the file.
-- file_edit_using_search_replace_blocks should start with <<<<<<< SEARCH
-""",
-        ),
-        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.",
-        ),
-        openai.pydantic_function_tool(
-            ContextSave,
-            description="""
-
-Saves provided description and file contents of all the relevant file paths or globs in a single text file.
-- Provide random unqiue id or whatever user provided.
-- Leave project path as empty string if no project path""",
-        ),
-    ]
-
-    system = initialize(
-        os.getcwd(),
-        [],
-        resume if (memory and resume) else "",
-        max_tokens=8000,
-        mode="wcgw",
-    )
-
-    with open(
-        os.path.join(
-            os.path.dirname(__file__), "..", "wcgw", "client", "diff-instructions.txt"
+            which_tool_name(tool.name), description=tool.description
         )
-    ) as f:
-        system += f.read()
-
-    if not history:
-        history = [{"role": "system", "content": system}]
-    else:
-        if history[-1]["role"] == "tool":
-            waiting_for_assistant = True
-
-    client = OpenAI()
+        for tool in TOOL_PROMPTS
+        if tool.name != "Initialize"
+    ]
 
     cost: float = 0
     input_toks = 0
@@ -262,203 +176,233 @@ Saves provided description and file contents of all the relevant file paths or g
     assistant_console = rich.console.Console(
         style="white bold", highlight=False, markup=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 waiting_for_assistant:
-            if first_message:
-                msg = first_message
-                first_message = ""
-            else:
-                msg = text_from_editor(user_console)
+    with BashState(
+        system_console, os.getcwd(), None, None, None, None, True, None
+    ) as bash_state:
+        context = Context(bash_state, system_console)
+        system, context = initialize(
+            "first_call",
+            context,
+            os.getcwd(),
+            [],
+            resume if (memory and resume) else "",
+            max_tokens=8000,
+            mode="wcgw",
+        )
 
-            history.append(parse_user_message_special(msg))
+        with open(
+            os.path.join(
+                os.path.dirname(__file__),
+                "..",
+                "wcgw",
+                "client",
+                "diff-instructions.txt",
+            )
+        ) as f:
+            system += f.read()
+
+        if not history:
+            history = [{"role": "system", "content": system}]
         else:
-            waiting_for_assistant = False
+            if history[-1]["role"] == "tool":
+                waiting_for_assistant = True
+
+        client = OpenAI()
+
+        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 waiting_for_assistant:
+                if first_message:
+                    msg = first_message
+                    first_message = ""
+                else:
+                    msg = text_from_editor(user_console)
+
+                history.append(parse_user_message_special(msg))
+            else:
+                waiting_for_assistant = 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,
-        )
+            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 = ""
-        image_histories: History = []
-        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": type(which_tool(tool_args)).__name__,
-                                },
-                            }
-                            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_dones, cost_ = get_tool_output(
-                                    json.loads(tool_args),
-                                    enc,
-                                    limit - cost,
-                                    loop,
-                                    max_tokens=8000,
-                                )
-                                output_or_done = output_or_dones[0]
-                            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(
+                "\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 = ""
+            image_histories: History = []
+            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": type(which_tool(tool_args)).__name__,
+                                    },
+                                }
+                                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_dones, cost_ = get_tool_output(
+                                        context,
+                                        json.loads(tool_args),
+                                        enc,
+                                        limit - cost,
+                                        loop,
+                                        max_tokens=8000,
+                                    )
+                                    output_or_done = output_or_dones[0]
+                                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}"
                                 )
-                                return output_or_done.task_output, cost
 
-                            output = output_or_done
-
-                            if isinstance(output, ImageData):
-                                randomId = petname.Generate(2, "-")
-                                if not image_histories:
-                                    image_histories.extend(
-                                        [
-                                            {
-                                                "role": "assistant",
-                                                "content": f"Share images with ids: {randomId}",
-                                            },
+                                output = output_or_done
+
+                                if isinstance(output, ImageData):
+                                    randomId = petname.Generate(2, "-")
+                                    if not image_histories:
+                                        image_histories.extend(
+                                            [
+                                                {
+                                                    "role": "assistant",
+                                                    "content": f"Share images with ids: {randomId}",
+                                                },
+                                                {
+                                                    "role": "user",
+                                                    "content": [
+                                                        {
+                                                            "type": "image_url",
+                                                            "image_url": {
+                                                                "url": output.dataurl,
+                                                                "detail": "auto",
+                                                            },
+                                                        }
+                                                    ],
+                                                },
+                                            ]
+                                        )
+                                    else:
+                                        image_histories[0]["content"] += ", " + randomId
+                                        second_content = image_histories[1]["content"]
+                                        assert isinstance(second_content, list)
+                                        second_content.append(
                                             {
-                                                "role": "user",
-                                                "content": [
-                                                    {
-                                                        "type": "image_url",
-                                                        "image_url": {
-                                                            "url": output.dataurl,
-                                                            "detail": "auto",
-                                                        },
-                                                    }
-                                                ],
-                                            },
-                                        ]
-                                    )
+                                                "type": "image_url",
+                                                "image_url": {
+                                                    "url": output.dataurl,
+                                                    "detail": "auto",
+                                                },
+                                            }
+                                        )
+
+                                    item = {
+                                        "role": "tool",
+                                        "content": f"Ask user for image id: {randomId}",
+                                        "tool_call_id": tool_call_id + str(toolindex),
+                                    }
                                 else:
-                                    image_histories[0]["content"] += ", " + randomId
-                                    second_content = image_histories[1]["content"]
-                                    assert isinstance(second_content, list)
-                                    second_content.append(
-                                        {
-                                            "type": "image_url",
-                                            "image_url": {
-                                                "url": output.dataurl,
-                                                "detail": "auto",
-                                            },
-                                        }
-                                    )
-
-                                item = {
-                                    "role": "tool",
-                                    "content": f"Ask user for image id: {randomId}",
-                                    "tool_call_id": tool_call_id + str(toolindex),
-                                }
-                            else:
-                                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)
-                    waiting_for_assistant = 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
+                                    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)
+                        waiting_for_assistant = 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_
 
-                chunk_str = chunk.choices[0].delta.content or ""
-                assistant_console.print(chunk_str, end="")
-                full_response += chunk_str
-        except KeyboardInterrupt:
-            waiting_for_assistant = False
-            input("Interrupted...enter to redo the current turn")
-        else:
-            history.extend(_histories)
-            history.extend(image_histories)
-            save_history(history, session_id)
+                        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:
+                waiting_for_assistant = False
+                input("Interrupted...enter to redo the current turn")
+            else:
+                history.extend(_histories)
+                history.extend(image_histories)
+                save_history(history, session_id)
 
     return "Couldn't finish the task", cost