wcgw 2.5.0__py3-none-any.whl → 2.6.2__py3-none-any.whl

wcgw/client/mcp_server/server.py

@@ -1,34 +1,33 @@
-import asyncio
 import importlib
 import json
 import os
-import sys
-import traceback
 from typing import Any
 
-from mcp_wcgw.server.models import InitializationOptions
+from pydantic import AnyUrl, ValidationError
+
+import mcp_wcgw.server.stdio
 import mcp_wcgw.types as types
-from mcp_wcgw.types import Tool as ToolParam
 from mcp_wcgw.server import NotificationOptions, Server
-from pydantic import AnyUrl, BaseModel, ValidationError
-import mcp_wcgw.server.stdio
-from .. import tools
-from ..tools import DoneFlag, get_tool_output, which_tool_name, default_enc
+from mcp_wcgw.server.models import InitializationOptions
+from mcp_wcgw.types import Tool as ToolParam
+
 from ...types_ import (
     BashCommand,
     BashInteraction,
-    WriteIfEmpty,
     FileEdit,
+    GetScreenInfo,
+    Initialize,
     Keyboard,
     Mouse,
     ReadFiles,
     ReadImage,
     ResetShell,
-    Initialize,
     ScreenShot,
-    GetScreenInfo,
+    WriteIfEmpty,
 )
+from .. import tools
 from ..computer_use import SLEEP_TIME_MAX_S
+from ..tools import DoneFlag, default_enc, get_tool_output, which_tool_name
 
 COMPUTER_USE_ON_DOCKER_ENABLED = False
 
@@ -77,6 +76,12 @@ async def handle_list_tools() -> list[types.Tool]:
             name="Initialize",
             description="""
 - Always call this at the start of the conversation before using any of the shell tools from wcgw.
+- This will reset the shell.
+- Use `any_workspace_path` to initialize the shell in the appropriate project directory.
+- If the user has mentioned a workspace or project root, use it to set `any_workspace_path`.
+- If the user has mentioned a folder or file with unclear project root, use the file or folder as `any_workspace_path`.
+- If user has mentioned any files use `initial_files_to_read` to read, use absolute paths only.
+- If `any_workspace_path` is provided, a tree structure of the workspace will be shown.
 """,
         ),
         ToolParam(
@@ -236,24 +241,24 @@ async def handle_call_tool(
         if isinstance(output_or_done, str):
             if issubclass(tool_type, Initialize):
                 output_or_done += """
-                
-    You're an expert software engineer with shell and code knowledge.
-
-    Instructions:
+---
+You're an expert software engineer with shell and code knowledge.
+
+Instructions:
+
+    - You should use the provided bash execution, reading and writing file tools to complete objective.
+    - First understand about the project by getting the folder structure (ignoring .git, node_modules, venv, etc.)
+    - Always read relevant files before editing.
+    - Do not provide code snippets unless asked by the user, instead directly add/edit the code.
+    - Do not install new tools/packages before ensuring no such tools/package or an alternative already exists.
+    - Do not use artifacts if you have access to the repository and not asked by the user to provide artifacts/snippets. Directly create/update using shell tools.
+    - Do not use Ctrl-c or Ctrl-z or interrupt commands without asking the user, because often the program don't show any update but they still are running.
+    - Do not use echo to write multi-line files, always use FileEdit tool to update a code.
     
-        - You should use the provided bash execution, reading and writing file tools to complete objective.
-        - First understand about the project by getting the folder structure (ignoring .git, node_modules, venv, etc.)
-        - Always read relevant files before editing.
-        - Do not provide code snippets unless asked by the user, instead directly add/edit the code.
-        - Do not install new tools/packages before ensuring no such tools/package or an alternative already exists.
-        - Do not use artifacts if you have access to the repository and not asked by the user to provide artifacts/snippets. Directly create/update using shell tools.
-        - Do not use Ctrl-c or Ctrl-z or interrupt commands without asking the user, because often the program don't show any update but they still are running.
-        - Do not use echo to write multi-line files, always use FileEdit tool to update a code.
-        
-    Additional instructions:
-        Always run `pwd` if you get any file or directory not found error to make sure you're not lost, or to get absolute cwd.
-
-        Always write production ready, syntactically correct code.
+Additional instructions:
+    Always run `pwd` if you get any file or directory not found error to make sure you're not lost, or to get absolute cwd.
+
+    Always write production ready, syntactically correct code.
     """
 
             content.append(types.TextContent(type="text", text=output_or_done))

wcgw/client/openai_client.py

@@ -1,58 +1,46 @@
 import base64
 import json
 import mimetypes
-from pathlib import Path
-import sys
+import os
+import subprocess
+import tempfile
 import traceback
-from typing import Callable, DefaultDict, Optional, cast
+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,
-    ChatCompletionAssistantMessageParam,
     ChatCompletionUserMessageParam,
-    ChatCompletionContentPartParam,
-    ChatCompletionMessage,
-    ParsedChatCompletionMessage,
 )
-import rich
-import petname  # type: ignore[import-untyped]
-import tokenizers  # type: ignore[import-untyped]
+from pydantic import BaseModel
 from typer import Typer
-import uuid
 
 from ..types_ import (
     BashCommand,
     BashInteraction,
-    WriteIfEmpty,
     FileEdit,
-    ReadImage,
     ReadFiles,
+    ReadImage,
     ResetShell,
+    WriteIfEmpty,
 )
-
-from .common import Models, discard_input
-from .common import CostData, History
+from .common import CostData, History, Models, discard_input
 from .openai_utils import get_input_cost, get_output_cost
-from .tools import ImageData
-
 from .tools import (
     DoneFlag,
+    ImageData,
     get_tool_output,
     which_tool,
 )
 
-from urllib import parse
-import subprocess
-import os
-import tempfile
-
-import toml
-from pydantic import BaseModel
-
-
-from dotenv import load_dotenv
-
 
 class Config(BaseModel):
     model: Models

wcgw/client/repo_ops/display_tree.py

@@ -0,0 +1,127 @@
+import io
+from pathlib import Path
+from typing import List, Set
+
+
+class DirectoryTree:
+    def __init__(self, root: Path, max_files: int = 10):
+        """
+        Initialize the DirectoryTree with a root path and maximum number of files to display
+
+        Args:
+            root_path: The root directory path to start from
+            max_files: Maximum number of files to display in unexpanded directories
+        """
+        self.root = root
+        self.max_files = max_files
+        self.expanded_files: Set[Path] = set()
+        self.expanded_dirs = set[Path]()
+
+        if not self.root.exists():
+            raise ValueError(f"Root path {root} does not exist")
+
+        if not self.root.is_dir():
+            raise ValueError(f"Root path {root} is not a directory")
+
+    def expand(self, rel_path: str) -> None:
+        """
+        Expand a specific file in the tree
+
+        Args:
+            rel_path: Relative path from root to the file to expand
+        """
+        abs_path = self.root / rel_path
+
+        if not abs_path.exists():
+            raise ValueError(f"Path {rel_path} does not exist")
+
+        if not abs_path.is_file():
+            raise ValueError(f"Path {rel_path} is not a file")
+
+        if not str(abs_path).startswith(str(self.root)):
+            raise ValueError(f"Path {rel_path} is outside root directory")
+
+        self.expanded_files.add(abs_path)
+
+        # Add all parent directories to expanded dirs
+        current = abs_path.parent
+        while str(current) >= str(self.root):
+            if current not in self.expanded_dirs:
+                self.expanded_dirs.add(current)
+            if current == current.parent:
+                break
+            current = current.parent
+
+    def _list_directory(self, dir_path: Path) -> List[Path]:
+        """List contents of a directory, sorted with directories first"""
+        contents = list(dir_path.iterdir())
+        return sorted(contents, key=lambda x: (not x.is_dir(), x.name.lower()))
+
+    def _count_hidden_items(
+        self, dir_path: Path, shown_items: List[Path]
+    ) -> tuple[int, int]:
+        """Count hidden files and directories in a directory"""
+        all_items = set(self._list_directory(dir_path))
+        shown_items_set = set(shown_items)
+        hidden_items = all_items - shown_items_set
+
+        hidden_files = sum(1 for p in hidden_items if p.is_file())
+        hidden_dirs = sum(1 for p in hidden_items if p.is_dir())
+
+        return hidden_files, hidden_dirs
+
+    def display(self) -> str:
+        """Display the directory tree with expanded state"""
+        writer = io.StringIO()
+
+        def _display_recursive(
+            current_path: Path, indent: int = 0, depth: int = 0
+        ) -> None:
+            # Print current directory name
+            if current_path == self.root:
+                writer.write(f"{current_path}\n")
+            else:
+                writer.write(f"{' ' * indent}{current_path.name}\n")
+
+            # Don't recurse beyond depth 1 unless path contains expanded files
+            if depth > 0 and current_path not in self.expanded_dirs:
+                return
+
+            # Get directory contents
+            contents = self._list_directory(current_path)
+            shown_items = []
+
+            for item in contents:
+                # Show items only if:
+                # 1. They are expanded files
+                # 2. They are parents of expanded items
+                should_show = item in self.expanded_files or item in self.expanded_dirs
+
+                if should_show:
+                    shown_items.append(item)
+                    if item.is_dir():
+                        _display_recursive(item, indent + 2, depth + 1)
+                    else:
+                        writer.write(f"{' ' * (indent + 2)}{item.name}\n")
+
+            # Show hidden items count if any items were hidden
+            hidden_files, hidden_dirs = self._count_hidden_items(
+                current_path, shown_items
+            )
+            if hidden_files > 0 or hidden_dirs > 0:
+                hidden_msg = []
+                if hidden_dirs > 0:
+                    hidden_msg.append(
+                        f"{hidden_dirs} director{'ies' if hidden_dirs != 1 else 'y'}"
+                    )
+                if hidden_files > 0:
+                    hidden_msg.append(
+                        f"{hidden_files} file{'s' if hidden_files != 1 else ''}"
+                    )
+                writer.write(
+                    f"{' ' * (indent + 2)}... {' and '.join(hidden_msg)} hidden\n"
+                )
+
+        _display_recursive(self.root, depth=0)
+
+        return writer.getvalue()

wcgw/client/repo_ops/path_prob.py

@@ -0,0 +1,58 @@
+from typing import Dict, List, Tuple
+
+import tokenizers  # type: ignore[import-untyped]
+
+
+class FastPathAnalyzer:
+    def __init__(self, model_path: str, vocab_path: str) -> None:
+        """Initialize with vocabulary."""
+        # Load vocabulary and probabilities
+        self.vocab_probs: Dict[str, float] = {}
+        with open(vocab_path, "r") as f:
+            for line in f:
+                parts = line.strip().split()
+                if len(parts) == 2:
+                    token, prob = parts
+                    try:
+                        self.vocab_probs[token] = float(prob)
+                    except ValueError:
+                        continue
+
+        self.encoder = tokenizers.Tokenizer.from_file(model_path)
+
+    def tokenize_batch(self, texts: List[str]) -> List[List[str]]:
+        """Tokenize multiple texts at once."""
+        encodings = self.encoder.encode_batch(texts)
+        return [encoding.tokens for encoding in encodings]
+
+    def detokenize(self, tokens: List[str]) -> str:
+        """Convert tokens back to text, handling special tokens."""
+        return self.encoder.decode(tokens)  # type: ignore[no-any-return]
+
+    def calculate_path_probabilities_batch(
+        self, paths: List[str]
+    ) -> List[Tuple[float, List[str], List[str]]]:
+        """Calculate log probability for multiple paths at once."""
+        # Batch tokenize all paths
+        all_tokens = self.tokenize_batch(paths)
+
+        results = []
+        for tokens in all_tokens:
+            # Calculate sum of log probabilities for each path
+            log_prob_sum = 0.0
+            unknown_tokens = []
+            for token in tokens:
+                if token in self.vocab_probs:
+                    log_prob_sum += self.vocab_probs[token]
+                else:
+                    unknown_tokens.append(token)
+
+            results.append((log_prob_sum, tokens, unknown_tokens))
+
+        return results
+
+    def calculate_path_probability(
+        self, path: str
+    ) -> Tuple[float, List[str], List[str]]:
+        """Calculate log probability for a single path."""
+        return self.calculate_path_probabilities_batch([path])[0]