wcgw 2.8.6__py3-none-any.whl → 2.8.7__py3-none-any.whl

wcgw/client/mcp_server/server.py

@@ -318,7 +318,8 @@ async def main(computer_use: bool) -> None:
     if computer_use:
         COMPUTER_USE_ON_DOCKER_ENABLED = True
 
-    version = importlib.metadata.version("wcgw")
+    version = str(importlib.metadata.version("wcgw"))
+    tools.console.log("wcgw version: " + version)
     # Run the server using stdin/stdout streams
     async with mcp_wcgw.server.stdio.stdio_server() as (read_stream, write_stream):
         await server.run(

{wcgw-2.8.6.dist-info → wcgw-2.8.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wcgw
-Version: 2.8.6
+Version: 2.8.7
 Summary: Shell and coding agent on claude and chatgpt
 Project-URL: Homepage, https://github.com/rusiaaman/wcgw
 Author-email: Aman Rusia <gapypi@arcfu.com>
@@ -28,12 +28,12 @@ Requires-Dist: websockets>=13.1
 Description-Content-Type: text/markdown
 
 # Shell and Coding agent for Claude and Chatgpt
+
 Empowering chat applications to code, build and run on your local machine.
 
 - Claude - An MCP server on claude desktop for autonomous shell and coding agent. (mac only)
 - Chatgpt - Allows custom gpt to talk to your shell via a relay server. (linux or mac)
 
-
 ⚠️ Warning: do not allow BashCommand tool without reviewing the command, it may result in data loss.
 
 [![Tests](https://github.com/rusiaaman/wcgw/actions/workflows/python-tests.yml/badge.svg?branch=main)](https://github.com/rusiaaman/wcgw/actions/workflows/python-tests.yml)
@@ -43,13 +43,14 @@ Empowering chat applications to code, build and run on your local machine.
 [![smithery badge](https://smithery.ai/badge/wcgw)](https://smithery.ai/server/wcgw)
 
 ## Updates
-- [15 Jan 2025] Modes introduced: architect, code-writer, and all powerful wcgw mode. 
+
+- [15 Jan 2025] Modes introduced: architect, code-writer, and all powerful wcgw mode.
 
 - [8 Jan 2025] Context saving tool for saving relevant file paths along with a description in a single file. Can be used as a task checkpoint or for knowledge transfer.
 
 - [29 Dec 2024] Syntax checking on file writing and edits is now stable. Made `initialize` tool call useful; sending smart repo structure to claude if any repo is referenced. Large file handling is also now improved.
 
-- [9 Dec 2024] [Vscode extension to paste context on Claude app](https://marketplace.visualstudio.com/items?itemName=AmanRusia.wcgw)  
+- [9 Dec 2024] [Vscode extension to paste context on Claude app](https://marketplace.visualstudio.com/items?itemName=AmanRusia.wcgw)
 
 - [01 Dec 2024] Removed author hosted relay server for chatgpt.
 
@@ -61,19 +62,19 @@ Empowering chat applications to code, build and run on your local machine.
 - ⚡ **Large file edit**: Supports large file incremental edits to avoid token limit issues. Faster than full file write.
 - ⚡ **Syntax checking on edits**: Reports feedback to the LLM if its edits have any syntax errors, so that it can redo it.
 - ⚡ **Interactive Command Handling**: Supports interactive commands using arrow keys, interrupt, and ansi escape sequences.
-- ⚡ **File protections**: 
-  - The AI needs to read a file at least once before it's allowed to edit or rewrite it. This avoids accidental overwrites. 
-  - Avoids context filling up while reading very large files. Files get chunked based on token length. 
+- ⚡ **File protections**:
+  - The AI needs to read a file at least once before it's allowed to edit or rewrite it. This avoids accidental overwrites.
+  - Avoids context filling up while reading very large files. Files get chunked based on token length.
   - On initialisation the provided workspace's directory structure is returned after selecting important files (based on .gitignore as well as a statistical approach)
   - File edit based on search-replace tries to find correct search block if it has multiple matches based on previous search blocks. Fails otherwise (for correctness).
   - File edit has spacing tolerant matching, with warning on issues like indentation mismatch. If there's no match, the closest match is returned to the AI to fix its mistakes.
   - Using Aider-like search and replace, which has better performance than tool call based search and replace.
-- ⚡ **Shell optimizations**: 
+- ⚡ **Shell optimizations**:
   - Only one command is allowed to be run at a time, simplifying management and avoiding rogue processes. There's only single shell instance at any point of time.
-  - Current working directory is always returned after any shell command to prevent AI from getting lost. 
+  - Current working directory is always returned after any shell command to prevent AI from getting lost.
   - Command polling exits after a quick timeout to avoid slow feedback. However, status checking has wait tolerance based on fresh output streaming from a command. Both of these approach combined provides a good shell interaction experience.
 - ⚡ **Saving repo context in a single file**: Task checkpointing using "ContextSave" tool saves detailed context in a single file. Tasks can later be resumed in a new chat asking "Resume `task id`". The saved file can be used to do other kinds of knowledge transfer, such as taking help from another AI.
-- ⚡ **Easily switch between various modes**: 
+- ⚡ **Easily switch between various modes**:
   - Ask it to run in 'architect' mode for planning. Inspired by adier's architect mode, work with Claude to come up with a plan first. Leads to better accuracy and prevents premature file editing.
   - Ask it to run in 'code-writer' mode for code editing and project building. You can provide specific paths with wild card support to prevent other files getting edited.
   - By default it runs in 'wcgw' mode that has no restrictions and full authorisation.
@@ -151,25 +152,29 @@ over here
 Then ask claude to execute shell commands, read files, edit files, run your code, etc.
 
 #### Task checkpoint or knowledge transfer
+
 - You can do a task checkpoint or a knowledge transfer by attaching "KnowledgeTransfer" prompt using "Attach from MCP" button.
 - On running "KnowledgeTransfer" prompt, the "ContextSave" tool will be called saving the task description and all file content together in a single file. An id for the task will be generated.
 - You can in a new chat say "Resume '<task id>'", the AI should then call "Initialize" with the task id and load the context from there.
 - Or you can directly open the file generated and share it with another AI for help.
 
 #### Modes
+
 There are three built-in modes. You may ask Claude to run in one of the modes, like "Use 'architect' mode"
-| **Mode**       | **Description**                                                             | **Allows**                                               | **Denies**                                    | **Invoke prompt**                                                                                  |
+| **Mode** | **Description** | **Allows** | **Denies** | **Invoke prompt** |
 |-----------------|-----------------------------------------------------------------------------|---------------------------------------------------------|----------------------------------------------|----------------------------------------------------------------------------------------------------|
-| **Architect**   | Designed for you to work with Claude to investigate and understand your repo. | Read-only commands                                       | FileEdit and Write tool                      | Run in mode='architect'                                                                          |
-| **Code-writer** | For code writing and development                                           | Specified path globs for editing or writing, specified commands | FileEdit for paths not matching specified glob, Write for paths not matching specified glob | Run in code writer mode, only 'tests/**' allowed, only uv command allowed |
-| **wcgw**        | Default mode with everything allowed                                       | Everything                                              | Nothing                                      | No prompt, or "Run in wcgw mode"                                                                                       |
+| **Architect** | Designed for you to work with Claude to investigate and understand your repo. | Read-only commands | FileEdit and Write tool | Run in mode='architect' |
+| **Code-writer** | For code writing and development | Specified path globs for editing or writing, specified commands | FileEdit for paths not matching specified glob, Write for paths not matching specified glob | Run in code writer mode, only 'tests/**' allowed, only uv command allowed |
+| **wcgw\*\* | Default mode with everything allowed | Everything | Nothing | No prompt, or "Run in wcgw mode" |
 
 Note: in code-writer mode either all commands are allowed or none are allowed for now. If you give a list of allowed commands, Claude is instructed to run only those commands, but no actual check happens. (WIP)
 
-### [Optional] Vs code extension 
+### [Optional] Vs code extension
+
 https://marketplace.visualstudio.com/items?itemName=AmanRusia.wcgw
 
-Commands: 
+Commands:
+
 - Select a text and press `cmd+'` and then enter instructions. This will switch the app to Claude and paste a text containing your instructions, file path, workspace dir, and the selected text.
 
 ## Chatgpt Setup
@@ -201,3 +206,36 @@ Then run
 `uvx --from wcgw@latest wcgw_local --claude`
 
 You can now directly write messages or press enter key to open vim for multiline message and text pasting.
+
+## Tools
+
+The server provides the following MCP tools:
+
+**Shell Operations:**
+
+- `Initialize`: Reset shell and set up workspace environment
+  - Parameters: `any_workspace_path` (string), `initial_files_to_read` (string[]), `mode_name` ("wcgw"|"architect"|"code_writer"), `task_id_to_resume` (string)
+- `BashCommand`: Execute shell commands with timeout control
+  - Parameters: `command` (string), `wait_for_seconds` (int, optional)
+- `BashInteraction`: Send keyboard input to running programs
+  - Parameters: `send_text` (string) or `send_specials` (["Enter"|"Key-up"|...]) or `send_ascii` (int[]), `wait_for_seconds` (int, optional)
+
+**File Operations:**
+
+- `ReadFiles`: Read content from one or more files
+  - Parameters: `file_paths` (string[])
+- `WriteIfEmpty`: Create new files or write to empty files
+  - Parameters: `file_path` (string), `file_content` (string)
+- `FileEdit`: Edit existing files using search/replace blocks
+  - Parameters: `file_path` (string), `file_edit_using_search_replace_blocks` (string)
+- `ReadImage`: Read image files for display/processing
+  - Parameters: `file_path` (string)
+
+**Project Management:**
+
+- `ContextSave`: Save project context and files for Knowledge Transfer or saving task checkpoints to be resumed later
+  - Parameters: `id` (string), `project_root_path` (string), `description` (string), `relevant_file_globs` (string[])
+- `ResetShell`: Emergency reset for shell environment
+  - Parameters: `should_reset` (boolean)
+
+All tools support absolute paths and include built-in protections against common errors. See the [MCP specification](https://modelcontextprotocol.io/) for detailed protocol information.

{wcgw-2.8.6.dist-info → wcgw-2.8.7.dist-info}/RECORD

@@ -12,7 +12,7 @@ wcgw/client/file_ops/diff_edit.py,sha256=OlJCpPSE_3T41q9H0yDORm6trjm3w6zh1EkuPTx
 wcgw/client/file_ops/search_replace.py,sha256=Napa7IWaYPGMNdttunKyRDkb90elZE7r23B_o_htRxo,5585
 wcgw/client/mcp_server/Readme.md,sha256=I8N4dHkTUVGNQ63BQkBMBhCCBTgqGOSF_pUR6iOEiUk,2495
 wcgw/client/mcp_server/__init__.py,sha256=hyPPwO9cabAJsOMWhKyat9yl7OlSmIobaoAZKHu3DMc,381
-wcgw/client/mcp_server/server.py,sha256=CBOS_DSDcBf5VlHQMJJbDL-wnn_DAPDvxX3dO8yodRo,13144
+wcgw/client/mcp_server/server.py,sha256=zMVa2nR2hapQZAtR34POch3VLVEEnJ6SQ8iWYLDxJrU,13199
 wcgw/client/repo_ops/display_tree.py,sha256=5FD4hfMkM2cIZnXlu7WfJswJLthj0SkuHlkGH6dpWQU,4632
 wcgw/client/repo_ops/path_prob.py,sha256=SWf0CDn37rtlsYRQ51ufSxay-heaQoVIhr1alB9tZ4M,2144
 wcgw/client/repo_ops/paths_model.vocab,sha256=M1pXycYDQehMXtpp-qAgU7rtzeBbCOiJo4qcYFY0kqk,315087
@@ -20,6 +20,12 @@ wcgw/client/repo_ops/paths_tokens.model,sha256=jiwwE4ae8ADKuTZISutXuM5Wfyc_FBmN5
 wcgw/client/repo_ops/repo_context.py,sha256=5NqRxBY0K-SBFXJ0Ybt7llzYOBD8pRkTpruMMJHWxv4,4336
 wcgw/relay/serve.py,sha256=Z5EwtaCAtKFBSnUw4mPYw0sze3Coc4Fa8gObRRG_bT0,9525
 wcgw/relay/static/privacy.txt,sha256=s9qBdbx2SexCpC_z33sg16TptmAwDEehMCLz4L50JLc,529
+wcgw_cli/__init__.py,sha256=TNxXsTPgb52OhakIda9wTRh91cqoBqgQRx5TxjzQQFU,21
+wcgw_cli/__main__.py,sha256=wcCrL4PjG51r5wVKqJhcoJPTLfHW0wNbD31DrUN0MWI,28
+wcgw_cli/anthropic_client.py,sha256=lZWEoX_qDOJIjzbG-EKxTjJyvTSw1Y5odtv7YUUIL7k,21054
+wcgw_cli/cli.py,sha256=GEje9ZBIaD5_-HK3zxZCGYaeDF8bfFxImloOR3O66Fw,1019
+wcgw_cli/openai_client.py,sha256=wp4XDf3t3W6XG5LHgr6bFckePyty24BGtsOEjOrIrk0,17955
+wcgw_cli/openai_utils.py,sha256=xGOb3W5ALrIozV7oszfGYztpj0FnXdD7jAxm5lEIVKY,2439
 mcp_wcgw/__init__.py,sha256=fKCgOdN7cn7gR3YGFaGyV5Goe8A2sEyllLcsRkN0i-g,2601
 mcp_wcgw/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mcp_wcgw/types.py,sha256=Enq5vqOPaQdObK9OQuafdLuMxb5RsLQ3k_k613fK41k,30556
@@ -42,8 +48,8 @@ mcp_wcgw/shared/memory.py,sha256=dBsOghxHz8-tycdSVo9kSujbsC8xb_tYsGmuJobuZnw,281
 mcp_wcgw/shared/progress.py,sha256=ymxOsb8XO5Mhlop7fRfdbmvPodANj7oq6O4dD0iUcnw,1048
 mcp_wcgw/shared/session.py,sha256=e44a0LQOW8gwdLs9_DE9oDsxqW2U8mXG3d5KT95bn5o,10393
 mcp_wcgw/shared/version.py,sha256=d2LZii-mgsPIxpshjkXnOTUmk98i0DT4ff8VpA_kAvE,111
-wcgw-2.8.6.dist-info/METADATA,sha256=PnCOOeODRmn7ZBHmoQN7qjS63YXOc0Xd78i7OVSo07o,12019
-wcgw-2.8.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-wcgw-2.8.6.dist-info/entry_points.txt,sha256=vd3tj1_Kzfp55LscJ8-6WFMM5hm9cWTfNGFCrWBnH3Q,124
-wcgw-2.8.6.dist-info/licenses/LICENSE,sha256=BvY8xqjOfc3X2qZpGpX3MZEmF-4Dp0LqgKBbT6L_8oI,11142
-wcgw-2.8.6.dist-info/RECORD,,
+wcgw-2.8.7.dist-info/METADATA,sha256=xXsXNlWtzj7OgpCgknOCaiiX1qG0KkOc8IbxsLW2k1A,13053
+wcgw-2.8.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+wcgw-2.8.7.dist-info/entry_points.txt,sha256=vd3tj1_Kzfp55LscJ8-6WFMM5hm9cWTfNGFCrWBnH3Q,124
+wcgw-2.8.7.dist-info/licenses/LICENSE,sha256=BvY8xqjOfc3X2qZpGpX3MZEmF-4Dp0LqgKBbT6L_8oI,11142
+wcgw-2.8.7.dist-info/RECORD,,

wcgw_cli/__init__.py

@@ -0,0 +1 @@
+from .cli import app

wcgw_cli/__main__.py

@@ -0,0 +1,3 @@
+from .cli import app
+
+app()

wcgw_cli/anthropic_client.py

@@ -0,0 +1,518 @@
+import base64
+import json
+import mimetypes
+import os
+import subprocess
+import tempfile
+import traceback
+import uuid
+from pathlib import Path
+from typing import Literal, Optional, cast
+
+import rich
+from anthropic import Anthropic
+from anthropic.types import (
+    ImageBlockParam,
+    MessageParam,
+    TextBlockParam,
+    ToolParam,
+    ToolResultBlockParam,
+    ToolUseBlockParam,
+)
+from dotenv import load_dotenv
+from typer import Typer
+
+from wcgw.client.common import discard_input
+from wcgw.client.memory import load_memory
+from wcgw.client.tools import (
+    DoneFlag,
+    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,
+)
+
+History = list[MessageParam]
+
+
+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)
+
+
+def parse_user_message_special(msg: str) -> MessageParam:
+    # Search for lines starting with `%` and treat them as special commands
+    parts: list[ImageBlockParam | TextBlockParam] = []
+    for line in msg.split("\n"):
+        if line.startswith("%"):
+            args = line[1:].strip().split(" ")
+            command = args[0]
+            assert command == "image"
+            image_path = " ".join(args[1:])
+            with open(image_path, "rb") as f:
+                image_bytes = f.read()
+                image_b64 = base64.b64encode(image_bytes).decode("utf-8")
+                image_type = mimetypes.guess_type(image_path)[0]
+            parts.append(
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": cast(
+                            'Literal["image/jpeg", "image/png", "image/gif", "image/webp"]',
+                            image_type or "image/png",
+                        ),
+                        "data": image_b64,
+                    },
+                }
+            )
+        else:
+            if len(parts) > 0 and parts[-1]["type"] == "text":
+                parts[-1]["text"] += "\n" + line
+            else:
+                parts.append({"type": "text", "text": line})
+    return {"role": "user", "content": parts}
+
+
+app = Typer(pretty_exceptions_show_locals=False)
+
+
+@app.command()
+def loop(
+    first_message: Optional[str] = None,
+    limit: Optional[float] = None,
+    resume: Optional[str] = None,
+    computer_use: bool = False,
+) -> tuple[str, float]:
+    load_dotenv()
+
+    session_id = str(uuid.uuid4())[:6]
+
+    history: History = []
+    waiting_for_assistant = False
+    memory = None
+    if resume:
+        try:
+            _, memory, _ = load_memory(
+                resume,
+                8000,
+                lambda x: default_enc.encode(x).ids,
+                lambda x: default_enc.decode(x),
+            )
+        except OSError:
+            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")
+            first_message = ""
+            waiting_for_assistant = history[-1]["role"] != "assistant"
+
+    limit = 1
+
+    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
+""",
+        ),
+    ]
+
+    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(
+        style="bright_black", highlight=False, markup=False
+    )
+    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)
+
+            history.append(parse_user_message_special(msg))
+        else:
+            waiting_for_assistant = False
+
+        cost_, input_toks_ = 0, 0
+        cost += cost_
+        input_toks += input_toks_
+
+        stream = client.messages.stream(
+            model="claude-3-5-sonnet-20241022",
+            messages=history,
+            tools=tools,
+            max_tokens=8096,
+            system=system,
+        )
+
+        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 type_ in {"message_start", "message_stop"}:
+                        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")
+                        ):
+                            chunk_str = content_block.text
+                            assistant_console.print(chunk_str, end="")
+                            full_response += chunk_str
+                        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 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
+                                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}"
+                            )
+
+                            _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,
+                                )
+                            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,
+                                            },
+                                        }
+                                    )
+
+                                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,
+                                )
+                            )
+                        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
+
+
+if __name__ == "__main__":
+    app()

wcgw_cli/cli.py

@@ -0,0 +1,42 @@
+import importlib
+from typing import Optional
+
+import typer
+from typer import Typer
+
+from wcgw_cli.anthropic_client import loop as claude_loop
+from wcgw_cli.openai_client import loop as openai_loop
+
+app = Typer(pretty_exceptions_show_locals=False)
+
+
+@app.command()
+def loop(
+    claude: bool = False,
+    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:
+        version_ = importlib.metadata.version("wcgw")
+        print(f"wcgw version: {version_}")
+        exit()
+    if claude:
+        return claude_loop(
+            first_message=first_message,
+            limit=limit,
+            resume=resume,
+            computer_use=computer_use,
+        )
+    else:
+        return openai_loop(
+            first_message=first_message,
+            limit=limit,
+            resume=resume,
+        )
+
+
+if __name__ == "__main__":
+    app()

wcgw_cli/openai_client.py

@@ -0,0 +1,467 @@
+import base64
+import json
+import mimetypes
+import os
+import subprocess
+import tempfile
+import traceback
+import uuid
+from pathlib import Path
+from typing import DefaultDict, Optional, cast
+
+import openai
+import petname  # type: ignore[import-untyped]
+import rich
+import tokenizers  # type: ignore[import-untyped]
+from dotenv import load_dotenv
+from openai import OpenAI
+from openai.types.chat import (
+    ChatCompletionContentPartParam,
+    ChatCompletionMessageParam,
+    ChatCompletionUserMessageParam,
+)
+from pydantic import BaseModel
+from typer import Typer
+
+from wcgw.client.common import CostData, History, Models, discard_input
+from wcgw.client.memory import load_memory
+from wcgw.client.tools import (
+    DoneFlag,
+    ImageData,
+    default_enc,
+    get_tool_output,
+    initialize,
+    which_tool,
+)
+from wcgw.types_ import (
+    BashCommand,
+    BashInteraction,
+    ContextSave,
+    FileEdit,
+    ReadFiles,
+    ReadImage,
+    ResetShell,
+    WriteIfEmpty,
+)
+
+from .openai_utils import get_input_cost, get_output_cost
+
+
+class Config(BaseModel):
+    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)
+
+
+def parse_user_message_special(msg: str) -> ChatCompletionUserMessageParam:
+    # Search for lines starting with `%` and treat them as special commands
+    parts: list[ChatCompletionContentPartParam] = []
+    for line in msg.split("\n"):
+        if line.startswith("%"):
+            args = line[1:].strip().split(" ")
+            command = args[0]
+            assert command == "image"
+            image_path = " ".join(args[1:])
+            with open(image_path, "rb") as f:
+                image_bytes = f.read()
+                image_b64 = base64.b64encode(image_bytes).decode("utf-8")
+                image_type = mimetypes.guess_type(image_path)[0]
+                dataurl = f"data:{image_type};base64,{image_b64}"
+            parts.append(
+                {"type": "image_url", "image_url": {"url": dataurl, "detail": "auto"}}
+            )
+        else:
+            if len(parts) > 0 and parts[-1]["type"] == "text":
+                parts[-1]["text"] += "\n" + line
+            else:
+                parts.append({"type": "text", "text": line})
+    return {"role": "user", "content": parts}
+
+
+app = Typer(pretty_exceptions_show_locals=False)
+
+
+@app.command()
+def loop(
+    first_message: Optional[str] = None,
+    limit: Optional[float] = None,
+    resume: Optional[str] = None,
+    computer_use: bool = False,
+) -> tuple[str, float]:
+    load_dotenv()
+
+    session_id = str(uuid.uuid4())[:6]
+
+    history: History = []
+    waiting_for_assistant = False
+
+    memory = None
+    if resume:
+        try:
+            _, memory, _ = load_memory(
+                resume,
+                8000,
+                lambda x: default_enc.encode(x).ids,
+                lambda x: default_enc.decode(x),
+            )
+        except OSError:
+            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")
+            first_message = ""
+            waiting_for_assistant = history[-1]["role"] != "assistant"
+
+    my_dir = os.path.dirname(__file__)
+
+    config = Config(
+        model=cast(Models, os.getenv("OPENAI_MODEL", "gpt-4o-2024-08-06").lower()),
+        cost_limit=0.1,
+        cost_unit="$",
+        cost_file={
+            "gpt-4o-2024-08-06": CostData(
+                cost_per_1m_input_tokens=5, cost_per_1m_output_tokens=15
+            ),
+        },
+    )
+
+    if limit is not None:
+        config.cost_limit = limit
+    limit = config.cost_limit
+
+    enc = tokenizers.Tokenizer.from_pretrained("Xenova/gpt-4o")
+
+    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"
+        )
+    ) 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()
+
+    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(
+        style="bright_black", highlight=False, markup=False
+    )
+    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)
+
+            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,
+        )
+
+        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(
+                                    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}",
+                                            },
+                                            {
+                                                "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(
+                                        {
+                                            "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
+                        )
+
+                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
+
+
+if __name__ == "__main__":
+    app()

wcgw_cli/openai_utils.py

@@ -0,0 +1,67 @@
+from typing import cast
+
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessage,
+    ChatCompletionMessageParam,
+    ParsedChatCompletionMessage,
+)
+from tokenizers import Tokenizer  # type: ignore[import-untyped]
+
+from wcgw.client.common import CostData, History
+
+
+def get_input_cost(
+    cost_map: CostData, enc: Tokenizer, history: History
+) -> tuple[float, int]:
+    input_tokens = 0
+    for msg in history:
+        content = msg["content"]
+        refusal = msg.get("refusal")
+        if isinstance(content, list):
+            for part in content:
+                if "text" in part:
+                    input_tokens += len(enc.encode(part["text"]))
+        elif content is None:
+            if refusal is None:
+                raise ValueError("Expected content or refusal to be present")
+            input_tokens += len(enc.encode(str(refusal)))
+        elif not isinstance(content, str):
+            raise ValueError(f"Expected content to be string, got {type(content)}")
+        else:
+            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: Tokenizer,
+    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