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

wcgw_cli/anthropic_client.py

@@ -10,11 +10,12 @@ from pathlib import Path
 from typing import Literal, Optional, cast
 
 import rich
-from anthropic import Anthropic
+from anthropic import Anthropic, MessageStopEvent
 from anthropic.types import (
     ImageBlockParam,
     MessageParam,
     ModelParam,
+    RawMessageStartEvent,
     TextBlockParam,
     ToolParam,
     ToolResultBlockParam,
@@ -24,30 +25,18 @@ from dotenv import load_dotenv
 from pydantic import BaseModel
 from typer import Typer
 
+from wcgw.client.bash_state.bash_state import BashState
 from wcgw.client.common import CostData, 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_name,
 )
-from wcgw.types_ import (
-    BashCommand,
-    BashInteraction,
-    ContextSave,
-    FileEdit,
-    GetScreenInfo,
-    Keyboard,
-    Mouse,
-    ReadFiles,
-    ReadImage,
-    ResetShell,
-    ScreenShot,
-    WriteIfEmpty,
-)
 
 
 class Config(BaseModel):
@@ -129,7 +118,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()
 
@@ -143,8 +131,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":
@@ -208,162 +196,14 @@ def loop(
 
     tools = [
         ToolParam(
-            input_schema=BashCommand.model_json_schema(),
-            name="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 "&".
-- Use longer wait_for_seconds if the command is expected to run for a long time.
-- Do not use 'cat' to read files, use ReadFiles tool instead.
-""",
-        ),
-        ToolParam(
-            input_schema=BashInteraction.model_json_schema(),
-            name="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.
-- This returns within 5 seconds, for heavy programs keep checking status for upto 10 turns before asking user to continue checking again.
-- Programs don't hang easily, so most likely explanation for no output is usually that the program is still running, and you need to check status again using ["Enter"].
-- Do not send Ctrl-c before checking for status till 10 minutes or whatever is appropriate for the program to finish.
-- Set longer wait_for_seconds when program is expected to run for a long time.
-""",
-        ),
-        ToolParam(
-            input_schema=ReadFiles.model_json_schema(),
-            name="ReadFiles",
-            description="""
-- Read full file content of one or more files.
-- Provide absolute file paths only
-""",
-        ),
-        ToolParam(
-            input_schema=WriteIfEmpty.model_json_schema(),
-            name="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.
-""",
-        ),
-        ToolParam(
-            input_schema=ReadImage.model_json_schema(),
-            name="ReadImage",
-            description="Read an image from the shell.",
-        ),
-        ToolParam(
-            input_schema=ResetShell.model_json_schema(),
-            name="ResetShell",
-            description="Resets the shell. Use only if all interrupts and prompt reset attempts have failed repeatedly.\nAlso exits the docker environment.\nYou need to call GetScreenInfo again",
-        ),
-        ToolParam(
-            input_schema=FileEdit.model_json_schema(),
-            name="FileEdit",
-            description="""
-- Use absolute file path only.
-- Use SEARCH/REPLACE blocks to edit the file.
-- If the edit fails due to block not matching, please retry with correct block till it matches. Re-read the file to ensure you've all the lines correct.
-""",
-        ),
-        ToolParam(
-            input_schema=ContextSave.model_json_schema(),
-            name="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
-""",
-        ),
+            name=tool.name,
+            description=tool.description,
+            input_schema=tool.inputSchema,
+        )
+        for tool in TOOL_PROMPTS
+        if tool.name != "Initialize"
     ]
 
-    if computer_use:
-        tools += [
-            ToolParam(
-                input_schema=GetScreenInfo.model_json_schema(),
-                name="GetScreenInfo",
-                description="""
-- Important: call this first in the conversation before ScreenShot, Mouse, and Keyboard tools.
-- Get display information of a linux os running on docker using image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
-- If user hasn't provided docker image id, check using `docker ps` and provide the id.
-- If the docker is not running, run using `docker run -d -p 6080:6080 ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest`
-- Connects shell to the docker environment.
-- Note: once this is called, the shell enters the docker environment. All bash commands will run over there.
-""",
-            ),
-            ToolParam(
-                input_schema=ScreenShot.model_json_schema(),
-                name="ScreenShot",
-                description="""
-- Capture screenshot of the linux os on docker.
-- All actions on UI using mouse and keyboard return within 0.5 seconds.
-    * So if you're doing something that takes longer for UI to update like heavy page loading, keep checking UI for update using ScreenShot upto 10 turns. 
-    * Notice for smallest of the loading icons to check if your action worked.
-    * After 10 turns of no change, ask user for permission to keep checking.
-    * If you don't notice even slightest of the change, it's likely you clicked on the wrong place.
-
-""",
-            ),
-            ToolParam(
-                input_schema=Mouse.model_json_schema(),
-                name="Mouse",
-                description="""
-- Interact with the linux os on docker using mouse.
-- Uses xdotool
-- About left_click_drag: the current mouse position will be used as the starting point, click and drag to the given x, y coordinates. Useful in things like sliders, moving things around, etc.
-- The output of this command has the screenshot after doing this action. Use this to verify if the action was successful.
-""",
-            ),
-            ToolParam(
-                input_schema=Keyboard.model_json_schema(),
-                name="Keyboard",
-                description="""
-- Interact with the linux os on docker using keyboard.
-- Emulate keyboard input to the screen
-- Uses xdootool to send keyboard input, keys like Return, BackSpace, Escape, Page_Up, etc. can be used.
-- Do not use it to interact with Bash tool.
-- Make sure you've selected a text area or an editable element before sending text.
-- The output of this command has the screenshot after doing this action. Use this to verify if the action was successful.
-""",
-            ),
-        ]
-
-    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"
-        )
-    ) as f:
-        system += f.read()
-
-    if history:
-        if (
-            (last_msg := history[-1])["role"] == "user"
-            and isinstance((content := last_msg["content"]), dict)
-            and content["type"] == "tool_result"
-        ):
-            waiting_for_assistant = True
-
-    client = Anthropic()
-
-    cost: float = 0
-    input_toks = 0
-    output_toks = 0
     system_console = rich.console.Console(style="blue", highlight=False, markup=False)
     error_console = rich.console.Console(style="red", highlight=False, markup=False)
     user_console = rich.console.Console(
@@ -372,216 +212,264 @@ 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: {config.cost_unit}{cost:.4f}, "
-                f"input tokens: {input_toks}"
-                f"output tokens: {output_toks}"
-            )
-            break
-        else:
-            system_console.print(
-                f"\nTotal cost: {config.cost_unit}{cost:.4f}, input tokens: {input_toks}, output tokens: {output_toks}"
+
+    with BashState(
+        system_console, os.getcwd(), None, None, None, None, False, None
+    ) as bash_state:
+        context = Context(bash_state, system_console)
+
+        system, context = initialize(
+            context,
+            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",
             )
+        ) as f:
+            system += f.read()
+
+        if history:
+            if (
+                (last_msg := history[-1])["role"] == "user"
+                and isinstance((content := last_msg["content"]), dict)
+                and content["type"] == "tool_result"
+            ):
+                waiting_for_assistant = True
 
-        if not waiting_for_assistant:
-            if first_message:
-                msg = first_message
-                first_message = ""
-            else:
-                msg = text_from_editor(user_console)
+        client = Anthropic()
 
-            history.append(parse_user_message_special(msg))
-        else:
-            waiting_for_assistant = False
-        stream = client.messages.stream(
-            model=config.model,
-            messages=history,
-            tools=tools,
-            max_tokens=8096,
-            system=system,
-        )
+        cost: float = 0
+        input_toks = 0
+        output_toks = 0
 
-        system_console.print(
-            "\n---------------------------------------\n# Assistant response",
-            style="bold",
-        )
-        _histories: History = []
-        full_response: str = ""
+        while True:
+            if cost > limit:
+                system_console.print(
+                    f"\nCost limit exceeded. Current cost: {config.cost_unit}{cost:.4f}, "
+                    f"input tokens: {input_toks}"
+                    f"output tokens: {output_toks}"
+                )
+                break
+            else:
+                system_console.print(
+                    f"\nTotal cost: {config.cost_unit}{cost:.4f}, input tokens: {input_toks}, output tokens: {output_toks}"
+                )
+
+            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
+
+            stream = client.messages.stream(
+                model=config.model,
+                messages=history,
+                tools=tools,
+                max_tokens=8096,
+                system=system,
+            )
 
-        tool_calls = []
-        tool_results: list[ToolResultBlockParam] = []
-        try:
-            with stream as stream_:
-                for chunk in stream_:
-                    type_ = chunk.type
-                    if type_ == "message_start":
-                        message_start = chunk.message
-                        # Update cost based on token usage from the API response
-                        input_tokens = message_start.usage.input_tokens
-                        input_toks += input_tokens
-                        cost += (
-                            input_tokens
-                            * config.cost_file[config.model].cost_per_1m_input_tokens
-                        ) / 1_000_000
-                    elif type_ == "message_stop":
-                        message_stop = chunk.message
-                        # Update cost based on output tokens
-                        output_tokens = message_stop.usage.output_tokens
-                        output_toks += output_tokens
-                        cost += (
-                            output_tokens
-                            * config.cost_file[config.model].cost_per_1m_output_tokens
-                        ) / 1_000_000
-                        continue
-                    elif type_ == "content_block_start" and hasattr(
-                        chunk, "content_block"
-                    ):
-                        content_block = chunk.content_block
-                        if (
-                            hasattr(content_block, "type")
-                            and content_block.type == "text"
-                            and hasattr(content_block, "text")
+            system_console.print(
+                "\n---------------------------------------\n# Assistant response",
+                style="bold",
+            )
+            _histories: History = []
+            full_response: str = ""
+
+            tool_calls = []
+            tool_results: list[ToolResultBlockParam] = []
+            try:
+                with stream as stream_:
+                    for chunk in stream_:
+                        type_ = chunk.type
+                        if isinstance(chunk, RawMessageStartEvent):
+                            message_start = chunk.message
+                            # Update cost based on token usage from the API response
+                            input_tokens = message_start.usage.input_tokens
+                            input_toks += input_tokens
+                            cost += (
+                                input_tokens
+                                * config.cost_file[
+                                    config.model
+                                ].cost_per_1m_input_tokens
+                            ) / 1_000_000
+                        elif isinstance(chunk, MessageStopEvent):
+                            message_stop = chunk.message
+                            # Update cost based on output tokens
+                            output_tokens = message_stop.usage.output_tokens
+                            output_toks += output_tokens
+                            cost += (
+                                output_tokens
+                                * config.cost_file[
+                                    config.model
+                                ].cost_per_1m_output_tokens
+                            ) / 1_000_000
+                            continue
+                        elif type_ == "content_block_start" and hasattr(
+                            chunk, "content_block"
                         ):
-                            chunk_str = content_block.text
-                            assistant_console.print(chunk_str, end="")
-                            full_response += chunk_str
-                        elif content_block.type == "tool_use":
+                            content_block = chunk.content_block
                             if (
-                                hasattr(content_block, "input")
-                                and hasattr(content_block, "name")
-                                and hasattr(content_block, "id")
+                                hasattr(content_block, "type")
+                                and content_block.type == "text"
+                                and hasattr(content_block, "text")
                             ):
-                                assert content_block.input == {}
-                                tool_calls.append(
-                                    {
-                                        "name": str(content_block.name),
-                                        "input": str(""),
-                                        "done": False,
-                                        "id": str(content_block.id),
-                                    }
-                                )
-                        else:
-                            error_console.log(
-                                f"Ignoring unknown content block type {content_block.type}"
-                            )
-                    elif type_ == "content_block_delta" and hasattr(chunk, "delta"):
-                        delta = chunk.delta
-                        if hasattr(delta, "type"):
-                            delta_type = str(delta.type)
-                            if delta_type == "text_delta" and hasattr(delta, "text"):
-                                chunk_str = delta.text
+                                chunk_str = content_block.text
                                 assistant_console.print(chunk_str, end="")
                                 full_response += chunk_str
-                            elif delta_type == "input_json_delta" and hasattr(
-                                delta, "partial_json"
-                            ):
-                                partial_json = delta.partial_json
-                                if isinstance(tool_calls[-1]["input"], str):
-                                    tool_calls[-1]["input"] += partial_json
+                            elif content_block.type == "tool_use":
+                                if (
+                                    hasattr(content_block, "input")
+                                    and hasattr(content_block, "name")
+                                    and hasattr(content_block, "id")
+                                ):
+                                    assert content_block.input == {}
+                                    tool_calls.append(
+                                        {
+                                            "name": str(content_block.name),
+                                            "input": str(""),
+                                            "done": False,
+                                            "id": str(content_block.id),
+                                        }
+                                    )
                             else:
                                 error_console.log(
-                                    f"Ignoring unknown content block delta type {delta_type}"
+                                    f"Ignoring unknown content block type {content_block.type}"
                                 )
-                        else:
-                            raise ValueError("Content block delta has no type")
-                    elif type_ == "content_block_stop":
-                        if tool_calls and not tool_calls[-1]["done"]:
-                            tc = tool_calls[-1]
-                            tool_name = str(tc["name"])
-                            tool_input = str(tc["input"])
-                            tool_id = str(tc["id"])
-
-                            tool_parsed = which_tool_name(
-                                tool_name
-                            ).model_validate_json(tool_input)
-
-                            system_console.print(
-                                f"\n---------------------------------------\n# Assistant invoked tool: {tool_parsed}"
-                            )
-
-                            _histories.append(
-                                {
-                                    "role": "assistant",
-                                    "content": [
-                                        ToolUseBlockParam(
-                                            id=tool_id,
-                                            name=tool_name,
-                                            input=tool_parsed.model_dump(),
-                                            type="tool_use",
-                                        )
-                                    ],
-                                }
-                            )
-                            try:
-                                output_or_dones, _ = get_tool_output(
-                                    tool_parsed,
-                                    default_enc,
-                                    limit - cost,
-                                    loop,
-                                    max_tokens=8000,
+                        elif type_ == "content_block_delta" and hasattr(chunk, "delta"):
+                            delta = chunk.delta
+                            if hasattr(delta, "type"):
+                                delta_type = str(delta.type)
+                                if delta_type == "text_delta" and hasattr(
+                                    delta, "text"
+                                ):
+                                    chunk_str = delta.text
+                                    assistant_console.print(chunk_str, end="")
+                                    full_response += chunk_str
+                                elif delta_type == "input_json_delta" and hasattr(
+                                    delta, "partial_json"
+                                ):
+                                    partial_json = delta.partial_json
+                                    if isinstance(tool_calls[-1]["input"], str):
+                                        tool_calls[-1]["input"] += partial_json
+                                else:
+                                    error_console.log(
+                                        f"Ignoring unknown content block delta type {delta_type}"
+                                    )
+                            else:
+                                raise ValueError("Content block delta has no type")
+                        elif type_ == "content_block_stop":
+                            if tool_calls and not tool_calls[-1]["done"]:
+                                tc = tool_calls[-1]
+                                tool_name = str(tc["name"])
+                                tool_input = str(tc["input"])
+                                tool_id = str(tc["id"])
+
+                                tool_parsed = which_tool_name(
+                                    tool_name
+                                ).model_validate_json(tool_input)
+
+                                system_console.print(
+                                    f"\n---------------------------------------\n# Assistant invoked tool: {tool_parsed}"
                                 )
-                            except Exception as e:
-                                output_or_dones = [
-                                    (f"GOT EXCEPTION while calling tool. Error: {e}")
-                                ]
-                                tb = traceback.format_exc()
-                                error_console.print(str(output_or_dones) + "\n" + tb)
-
-                            if any(isinstance(x, DoneFlag) for x in output_or_dones):
-                                return "", cost
-
-                            tool_results_content: list[
-                                TextBlockParam | ImageBlockParam
-                            ] = []
-                            for output in output_or_dones:
-                                assert not isinstance(output, DoneFlag)
-                                if isinstance(output, ImageData):
-                                    tool_results_content.append(
-                                        {
-                                            "type": "image",
-                                            "source": {
-                                                "type": "base64",
-                                                "media_type": output.media_type,
-                                                "data": output.data,
-                                            },
-                                        }
+
+                                _histories.append(
+                                    {
+                                        "role": "assistant",
+                                        "content": [
+                                            ToolUseBlockParam(
+                                                id=tool_id,
+                                                name=tool_name,
+                                                input=tool_parsed.model_dump(),
+                                                type="tool_use",
+                                            )
+                                        ],
+                                    }
+                                )
+                                try:
+                                    output_or_dones, _ = get_tool_output(
+                                        context,
+                                        tool_parsed,
+                                        default_enc,
+                                        limit - cost,
+                                        loop,
+                                        max_tokens=8000,
+                                    )
+                                except Exception as e:
+                                    output_or_dones = [
+                                        (
+                                            f"GOT EXCEPTION while calling tool. Error: {e}"
+                                        )
+                                    ]
+                                    tb = traceback.format_exc()
+                                    error_console.print(
+                                        str(output_or_dones) + "\n" + tb
                                     )
 
-                                else:
-                                    tool_results_content.append(
-                                        {
-                                            "type": "text",
-                                            "text": output,
-                                        },
+                                tool_results_content: list[
+                                    TextBlockParam | ImageBlockParam
+                                ] = []
+                                for output in output_or_dones:
+                                    if isinstance(output, ImageData):
+                                        tool_results_content.append(
+                                            {
+                                                "type": "image",
+                                                "source": {
+                                                    "type": "base64",
+                                                    "media_type": output.media_type,
+                                                    "data": output.data,
+                                                },
+                                            }
+                                        )
+
+                                    else:
+                                        tool_results_content.append(
+                                            {
+                                                "type": "text",
+                                                "text": output,
+                                            },
+                                        )
+                                tool_results.append(
+                                    ToolResultBlockParam(
+                                        type="tool_result",
+                                        tool_use_id=str(tc["id"]),
+                                        content=tool_results_content,
                                     )
-                            tool_results.append(
-                                ToolResultBlockParam(
-                                    type="tool_result",
-                                    tool_use_id=str(tc["id"]),
-                                    content=tool_results_content,
                                 )
-                            )
-                        else:
-                            _histories.append(
-                                {
-                                    "role": "assistant",
-                                    "content": full_response
-                                    if full_response.strip()
-                                    else "...",
-                                }  # Fixes anthropic issue of non empty response only
-                            )
-
-        except KeyboardInterrupt:
-            waiting_for_assistant = False
-            input("Interrupted...enter to redo the current turn")
-        else:
-            history.extend(_histories)
-            if tool_results:
-                history.append({"role": "user", "content": tool_results})
-                waiting_for_assistant = True
-            save_history(history, session_id)
+                            else:
+                                _histories.append(
+                                    {
+                                        "role": "assistant",
+                                        "content": full_response
+                                        if full_response.strip()
+                                        else "...",
+                                    }  # Fixes anthropic issue of non empty response only
+                                )
+
+            except KeyboardInterrupt:
+                waiting_for_assistant = False
+                input("Interrupted...enter to redo the current turn")
+            else:
+                history.extend(_histories)
+                if tool_results:
+                    history.append({"role": "user", "content": tool_results})
+                    waiting_for_assistant = True
+                save_history(history, session_id)
 
     return "Couldn't finish the task", cost