wcgw 2.8.7__tar.gz → 2.8.9__tar.gz

{wcgw-2.8.7 → wcgw-2.8.9}/PKG-INFO

@@ -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 → wcgw-2.8.9}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 authors = [{ name = "Aman Rusia", email = "gapypi@arcfu.com" }]
 name = "wcgw"
-version = "2.8.7"
+version = "2.8.9"
 description = "Shell and coding agent on claude and chatgpt"
 readme = "README.md"
 requires-python = ">=3.11, <3.13"

{wcgw-2.8.7 → wcgw-2.8.9}/src/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 → wcgw-2.8.9}/src/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"

{wcgw-2.8.7 → wcgw-2.8.9}/tests/test_initialize.py

@@ -248,3 +248,78 @@ class TestInitialize(unittest.TestCase):
 
             # Verify task memory was still loaded despite state error
             self.assertIn("Following is the retrieved task:\ntest_memory", result)
+
+
+    def test_workspace_path_is_file(self):
+        """Test initialize when workspace path points to a file"""
+        # Create a test file
+        test_file = os.path.join(self.test_workspace, "test.py")
+        os.makedirs(os.path.dirname(test_file), exist_ok=True)
+        with open(test_file, "w") as f:
+            f.write("print('test')")
+
+        with (
+            patch("os.path.exists") as mock_exists,
+            patch("os.path.isfile") as mock_isfile,
+            patch("wcgw.client.tools.get_repo_context") as mock_get_context,
+            patch("wcgw.client.tools.read_files") as mock_read_files,
+        ):
+            mock_exists.return_value = True
+            mock_isfile.return_value = True
+            mock_get_context.return_value = (self.repo_context, os.path.dirname(test_file))
+            mock_read_files.return_value = f"``` {test_file}print('test')```"
+
+            # Call initialize with file path and no read_files_
+            result = initialize(
+                any_workspace_path=test_file,
+                read_files_=[],
+                task_id_to_resume="",
+                max_tokens=None,
+                mode=Modes.wcgw,
+            )
+
+            # Verify read_files was called with the file path
+            mock_read_files.assert_called_once_with([test_file], None)
+
+            # Verify repo context uses parent directory
+            mock_get_context.assert_called_once_with(os.path.dirname(test_file), 200)
+
+            # Verify file content is in output
+            self.assertIn(f"``` {test_file}print('test')", result)
+
+            # Verify cwd
+            self.assertIn(f"Initialized in directory (also cwd): {os.path.dirname(test_file)}\n", result)
+
+    def test_workspace_path_is_file_with_read_files(self):
+        """Test initialize when workspace path points to a file and read_files_ is provided"""
+        test_file = os.path.join(self.test_workspace, "test.py")
+        extra_file = os.path.join(self.test_workspace, "extra.py")
+
+        with (
+            patch("os.path.exists") as mock_exists,
+            patch("os.path.isfile") as mock_isfile,
+            patch("wcgw.client.tools.get_repo_context") as mock_get_context,
+            patch("wcgw.client.tools.read_files") as mock_read_files,
+        ):
+            mock_exists.return_value = True
+            mock_isfile.return_value = True
+            mock_get_context.return_value = (self.repo_context, os.path.dirname(test_file))
+            mock_read_files.return_value = f"``` {extra_file}extra content```"
+
+            # Call initialize with file path and explicit read_files_
+            result = initialize(
+                any_workspace_path=test_file,
+                read_files_=[extra_file],
+                task_id_to_resume="",
+                max_tokens=None,
+                mode=Modes.wcgw,
+            )
+
+            # Verify read_files was called with the provided read_files_ only
+            mock_read_files.assert_called_once_with([extra_file], None)
+
+            # Verify repo context uses parent directory
+            mock_get_context.assert_called_once_with(os.path.dirname(test_file), 200)
+
+            # Verify extra file content is in output
+            self.assertIn(f"``` {extra_file}extra content", result)