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

wcgw/client/tools.py

@@ -146,6 +146,7 @@ def start_shell(is_restricted_mode: bool, initial_dir: str) -> pexpect.spawn:  #
             encoding="utf-8",
             timeout=TIMEOUT,
             cwd=initial_dir,
+            codec_errors="backslashreplace",
         )
         shell.sendline(
             f"export PROMPT_COMMAND= PS1={PROMPT_CONST}"
@@ -161,6 +162,7 @@ def start_shell(is_restricted_mode: bool, initial_dir: str) -> pexpect.spawn:  #
             echo=False,
             encoding="utf-8",
             timeout=TIMEOUT,
+            codec_errors="backslashreplace",
         )
         shell.sendline(f"export PS1={PROMPT_CONST}")
         shell.expect(PROMPT_CONST, timeout=TIMEOUT)
@@ -256,7 +258,9 @@ class BashState:
             before = "\n".join(before_lines).strip()
             counts += 1
             if counts > 100:
-                raise ValueError("Error in understanding shell output. This shouldn't happen, likely shell is in a bad state, please reset it")
+                raise ValueError(
+                    "Error in understanding shell output. This shouldn't happen, likely shell is in a bad state, please reset it"
+                )
 
         try:
             return int(before)
@@ -273,7 +277,7 @@ class BashState:
             self._bash_command_mode.bash_mode == "restricted_mode",
             self._cwd,
         )
-    
+
         self._pending_output = ""
 
         # Get exit info to ensure shell is ready
@@ -414,7 +418,9 @@ class BashState:
                 index = self.shell.expect([self._prompt, pexpect.TIMEOUT], timeout=0.2)
                 counts += 1
                 if counts > 100:
-                    raise ValueError("Error in understanding shell output. This shouldn't happen, likely shell is in a bad state, please reset it")
+                    raise ValueError(
+                        "Error in understanding shell output. This shouldn't happen, likely shell is in a bad state, please reset it"
+                    )
             console.print(f"Prompt updated to: {self._prompt}")
             return True
         return False
@@ -457,6 +463,12 @@ def initialize(
     folder_to_start = None
     if any_workspace_path:
         if os.path.exists(any_workspace_path):
+            if os.path.isfile(any_workspace_path):
+                # Set any_workspace_path to the directory containing the file
+                # Add the file to read_files_ only if empty to avoid duplicates
+                if not read_files_:
+                    read_files_ = [any_workspace_path]
+                any_workspace_path = os.path.dirname(any_workspace_path)
             repo_context, folder_to_start = get_repo_context(any_workspace_path, 200)
 
             repo_context = f"---\n# Workspace structure\n{repo_context}\n---\n"

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wcgw
-Version: 2.8.7
+Version: 2.8.9
 Summary: Shell and coding agent on claude and chatgpt
 Project-URL: Homepage, https://github.com/rusiaaman/wcgw
 Author-email: Aman Rusia <gapypi@arcfu.com>

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

@@ -7,7 +7,7 @@ wcgw/client/diff-instructions.txt,sha256=tmJ9Fu9XdO_72lYXQQNY9RZyx91bjxrXJf9d_KB
 wcgw/client/memory.py,sha256=8LdYsOhvCOoC1kfvDr85kNy07WnhPMvE6B2FRM2w85Y,2902
 wcgw/client/modes.py,sha256=FkDJIgjKrlJEufLq3abWfqV25BdF2pH-HnoHafy9LrA,10484
 wcgw/client/sys_utils.py,sha256=GajPntKhaTUMn6EOmopENWZNR2G_BJyuVbuot0x6veI,1376
-wcgw/client/tools.py,sha256=ndIvlJlULCxqRlks9cfDJK8JqMK0Ry85Mt4kZQ0IZhQ,52216
+wcgw/client/tools.py,sha256=xgLsgkWZ5qRiLpdyF_1oKIJGy-tj9_WmJQW30-6KhhQ,52759
 wcgw/client/file_ops/diff_edit.py,sha256=OlJCpPSE_3T41q9H0yDORm6trjm3w6zh1EkuPTxik2A,16832
 wcgw/client/file_ops/search_replace.py,sha256=Napa7IWaYPGMNdttunKyRDkb90elZE7r23B_o_htRxo,5585
 wcgw/client/mcp_server/Readme.md,sha256=I8N4dHkTUVGNQ63BQkBMBhCCBTgqGOSF_pUR6iOEiUk,2495
@@ -22,7 +22,7 @@ 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/anthropic_client.py,sha256=ZMFHD6g6h_PAReL8tubI8LnxVeRA3hcFLGNitjt9XhQ,24047
 wcgw_cli/cli.py,sha256=GEje9ZBIaD5_-HK3zxZCGYaeDF8bfFxImloOR3O66Fw,1019
 wcgw_cli/openai_client.py,sha256=wp4XDf3t3W6XG5LHgr6bFckePyty24BGtsOEjOrIrk0,17955
 wcgw_cli/openai_utils.py,sha256=xGOb3W5ALrIozV7oszfGYztpj0FnXdD7jAxm5lEIVKY,2439
@@ -48,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.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-2.8.9.dist-info/METADATA,sha256=5QkVxDqkY5Flyvrua5JMbLoKnbdDlXfr4UdiLAJ3Ad0,13053
+wcgw-2.8.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+wcgw-2.8.9.dist-info/entry_points.txt,sha256=vd3tj1_Kzfp55LscJ8-6WFMM5hm9cWTfNGFCrWBnH3Q,124
+wcgw-2.8.9.dist-info/licenses/LICENSE,sha256=BvY8xqjOfc3X2qZpGpX3MZEmF-4Dp0LqgKBbT6L_8oI,11142
+wcgw-2.8.9.dist-info/RECORD,,

wcgw_cli/anthropic_client.py

@@ -14,15 +14,17 @@ from anthropic import Anthropic
 from anthropic.types import (
     ImageBlockParam,
     MessageParam,
+    ModelParam,
     TextBlockParam,
     ToolParam,
     ToolResultBlockParam,
     ToolUseBlockParam,
 )
 from dotenv import load_dotenv
+from pydantic import BaseModel
 from typer import Typer
 
-from wcgw.client.common import discard_input
+from wcgw.client.common import CostData, discard_input
 from wcgw.client.memory import load_memory
 from wcgw.client.tools import (
     DoneFlag,
@@ -47,6 +49,14 @@ from wcgw.types_ import (
     WriteIfEmpty,
 )
 
+
+class Config(BaseModel):
+    model: ModelParam
+    cost_limit: float
+    cost_file: dict[ModelParam, CostData]
+    cost_unit: str = "$"
+
+
 History = list[MessageParam]
 
 
@@ -150,7 +160,51 @@ def loop(
             first_message = ""
             waiting_for_assistant = history[-1]["role"] != "assistant"
 
-    limit = 1
+    config = Config(
+        model="claude-3-5-sonnet-20241022",
+        cost_limit=0.1,
+        cost_unit="$",
+        cost_file={
+            # Claude 3.5 Haiku
+            "claude-3-5-haiku-latest": CostData(
+                cost_per_1m_input_tokens=0.80, cost_per_1m_output_tokens=4
+            ),
+            "claude-3-5-haiku-20241022": CostData(
+                cost_per_1m_input_tokens=0.80, cost_per_1m_output_tokens=4
+            ),
+            # Claude 3.5 Sonnet
+            "claude-3-5-sonnet-latest": CostData(
+                cost_per_1m_input_tokens=3.0, cost_per_1m_output_tokens=15.0
+            ),
+            "claude-3-5-sonnet-20241022": CostData(
+                cost_per_1m_input_tokens=3.0, cost_per_1m_output_tokens=15.0
+            ),
+            "claude-3-5-sonnet-20240620": CostData(
+                cost_per_1m_input_tokens=3.0, cost_per_1m_output_tokens=15.0
+            ),
+            # Claude 3 Opus
+            "claude-3-opus-latest": CostData(
+                cost_per_1m_input_tokens=15.0, cost_per_1m_output_tokens=75.0
+            ),
+            "claude-3-opus-20240229": CostData(
+                cost_per_1m_input_tokens=15.0, cost_per_1m_output_tokens=75.0
+            ),
+            # Legacy Models
+            "claude-3-haiku-20240307": CostData(
+                cost_per_1m_input_tokens=0.25, cost_per_1m_output_tokens=1.25
+            ),
+            "claude-2.1": CostData(
+                cost_per_1m_input_tokens=8.0, cost_per_1m_output_tokens=24.0
+            ),
+            "claude-2.0": CostData(
+                cost_per_1m_input_tokens=8.0, cost_per_1m_output_tokens=24.0
+            ),
+        },
+    )
+
+    if limit is not None:
+        config.cost_limit = limit
+    limit = config.cost_limit
 
     tools = [
         ToolParam(
@@ -321,9 +375,15 @@ Saves provided description and file contents of all the relevant file paths or g
     while True:
         if cost > limit:
             system_console.print(
-                f"\nCost limit exceeded. Current cost: {cost}, input tokens: {input_toks}, output tokens: {output_toks}"
+                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:
@@ -335,13 +395,8 @@ Saves provided description and file contents of all the relevant file paths or g
             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",
+            model=config.model,
             messages=history,
             tools=tools,
             max_tokens=8096,
@@ -361,7 +416,24 @@ Saves provided description and file contents of all the relevant file paths or g
             with stream as stream_:
                 for chunk in stream_:
                     type_ = chunk.type
-                    if type_ in {"message_start", "message_stop"}:
+                    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"