wcgw 2.6.3__py3-none-any.whl → 2.7.1__py3-none-any.whl

wcgw/client/file_ops/search_replace.py

@@ -0,0 +1,119 @@
+import re
+from typing import Callable
+
+from .diff_edit import FileEditInput, FileEditOutput
+
+
+def search_replace_edit(
+    lines: list[str], original_content: str, logger: Callable[[str], object]
+) -> tuple[str, str]:
+    if not lines:
+        raise Exception("Error: No input to search replace edit")
+    original_lines = original_content.split("\n")
+    n_lines = len(lines)
+    i = 0
+    search_replace_blocks = list[tuple[list[str], list[str]]]()
+    while i < n_lines:
+        if re.match(r"^<<<<<<+\s*SEARCH\s*$", lines[i]):
+            search_block = []
+            i += 1
+            while i < n_lines and not re.match(r"^======*\s*$", lines[i]):
+                search_block.append(lines[i])
+                i += 1
+            i += 1
+            if not search_block:
+                raise Exception("SEARCH block can not be empty")
+            replace_block = []
+            while i < n_lines and not re.match(r"^>>>>>>+\s*REPLACE\s*$", lines[i]):
+                replace_block.append(lines[i])
+                i += 1
+            i += 1
+
+            for line in search_block:
+                logger("> " + line)
+            logger("=======")
+            for line in replace_block:
+                logger("< " + line)
+            logger("\n\n\n\n")
+
+            search_replace_blocks.append((search_block, replace_block))
+        else:
+            i += 1
+
+    if not search_replace_blocks:
+        raise Exception(
+            "No valid search replace blocks found, ensure your SEARCH/REPLACE blocks are formatted correctly"
+        )
+
+    edited_content, comments_ = greedy_context_replace(
+        original_lines, [[x] for x in search_replace_blocks], original_lines, set(), 0
+    )
+    edited_file = "\n".join(edited_content)
+    if not comments_:
+        comments = "Edited successfully"
+    else:
+        comments = (
+            "Edited successfully. However, following warnings were generated while matching search blocks.\n"
+            + "\n".join(comments_)
+        )
+    return edited_file, comments
+
+
+def greedy_context_replace(
+    original_lines: list[str],
+    search_replace_blocks: list[list[tuple[list[str], list[str]]]],
+    running_lines: list[str],
+    running_comments: set[str],
+    current_block_offset: int,
+) -> tuple[list[str], set[str]]:
+    if current_block_offset >= len(search_replace_blocks):
+        return running_lines, running_comments
+    current_blocks = search_replace_blocks[current_block_offset]
+
+    outputs = FileEditInput(running_lines, 0, current_blocks, 0).edit_file()
+    best_matches, is_error = FileEditOutput.get_best_match(outputs)
+
+    if is_error:
+        best_matches[0].replace_or_throw(3)
+        raise Exception("Shouldn't happen")
+
+    if len(best_matches) > 1:
+        # Duplicate found, try to ground using previous blocks.
+        if current_block_offset == 0:
+            raise Exception(f"""
+    The following block matched more than once:
+    ---
+    ```
+    {'\n'.join(current_blocks[-1][0])}
+    ```
+    """)
+
+        else:
+            search_replace_blocks = (
+                search_replace_blocks[: current_block_offset - 1]
+                + [search_replace_blocks[current_block_offset - 1] + current_blocks]
+                + search_replace_blocks[current_block_offset + 1 :]
+            )
+            try:
+                return greedy_context_replace(
+                    original_lines, search_replace_blocks, original_lines, set(), 0
+                )
+            except Exception:
+                raise Exception(f"""
+        The following block matched more than once:
+        ---
+        ```
+        {'\n'.join(current_blocks[-1][0])}
+        ```
+        """)
+
+    best_match = best_matches[0]
+    running_lines, comments = best_match.replace_or_throw(3)
+    running_comments = running_comments | comments
+    return greedy_context_replace(
+        original_lines,
+        search_replace_blocks,
+        running_lines,
+        running_comments,
+        current_block_offset + 1,
+    )

wcgw/client/mcp_server/server.py

@@ -14,6 +14,7 @@ from mcp_wcgw.types import Tool as ToolParam
 from ...types_ import (
     BashCommand,
     BashInteraction,
+    ContextSave,
     FileEdit,
     GetScreenInfo,
     Initialize,
@@ -46,14 +47,47 @@ async def handle_read_resource(uri: AnyUrl) -> str:
 
 @server.list_prompts()  # type: ignore
 async def handle_list_prompts() -> list[types.Prompt]:
-    return []
+    return [
+        types.Prompt(
+            name="KnowledgeTransfer",
+            description="Prompt for invoking ContextSave tool in order to do a comprehensive knowledge transfer of a coding task. Prompts to save detailed error log and instructions.",
+        )
+    ]
 
 
 @server.get_prompt()  # type: ignore
 async def handle_get_prompt(
     name: str, arguments: dict[str, str] | None
 ) -> types.GetPromptResult:
-    return types.GetPromptResult(messages=[])
+    messages = []
+    if name == "KnowledgeTransfer":
+        messages = [
+            types.PromptMessage(
+                role="user",
+                content=types.TextContent(
+                    type="text",
+                    text="""Use `ContextSave` tool to do a knowledge transfer of the task in hand.
+Write detailed description in order to do a KT.
+Save all information necessary for a person to understand the task and the problems.
+
+Format the `description` field using Markdown with the following sections.
+- "# Objective" section containing project and task objective.
+- "# All user instructions" section should be provided containing all instructions user shared in the conversation.
+- "# Current status of the task" should be provided containing only what is already achieved, not what's remaining.
+- "# All issues with snippets" section containing snippets of error, traceback, file snippets, commands, etc. But no comments or solutions.
+- Be very verbose in the all issues with snippets section providing as much error context as possible.
+- "# Build and development instructions" section containing instructions to build or run project or run tests, or envrionment related information. Only include what's known. Leave empty if unknown.
+- After the tool completes succesfully, tell me the task id and the file path the tool generated (important!)
+- This tool marks end of your conversation, do not run any further tools after calling this.
+
+Provide all relevant file paths in order to understand and solve the the task. Err towards providing more file paths than fewer.
+
+(Note to self: this conversation can then be resumed later asking "Resume `<generated id>`" which should call Initialize tool)
+""",
+                ),
+            )
+        ]
+    return types.GetPromptResult(messages=messages)
 
 
 @server.list_tools()  # type: ignore
@@ -82,6 +116,7 @@ async def handle_list_tools() -> list[types.Tool]:
 - 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.
+- Leave `any_workspace_path` as empty if no file or folder is mentioned.
 """,
         ),
         ToolParam(
@@ -152,6 +187,14 @@ async def handle_list_tools() -> list[types.Tool]:
 """
             + diffinstructions,
         ),
+        ToolParam(
+            inputSchema=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_ON_DOCKER_ENABLED:
         tools += [

wcgw/client/memory.py

@@ -0,0 +1,78 @@
+import os
+import re
+import shlex
+from typing import Callable, Optional
+
+from ..types_ import ContextSave
+
+
+def get_app_dir_xdg() -> str:
+    xdg_data_dir = os.environ.get("XDG_DATA_HOME", os.path.expanduser("~/.local/share"))
+    return os.path.join(xdg_data_dir, "wcgw")
+
+
+def format_memory(task_memory: ContextSave, relevant_files: str) -> str:
+    memory_data = ""
+    if task_memory.project_root_path:
+        memory_data += (
+            f"# PROJECT ROOT = {shlex.quote(task_memory.project_root_path)}\n"
+        )
+    memory_data += task_memory.description
+
+    memory_data += (
+        "\n\n"
+        + "# Relevant file paths\n"
+        + ", ".join(map(shlex.quote, task_memory.relevant_file_globs))
+    )
+
+    memory_data += "\n\n# Relevant Files:\n" + relevant_files
+
+    return memory_data
+
+
+def save_memory(task_memory: ContextSave, relevant_files: str) -> str:
+    app_dir = get_app_dir_xdg()
+    memory_dir = os.path.join(app_dir, "memory")
+    os.makedirs(memory_dir, exist_ok=True)
+
+    task_id = task_memory.id
+    if not task_id:
+        raise Exception("Task id can not be empty")
+    memory_data = format_memory(task_memory, relevant_files)
+
+    memory_file_full = os.path.join(memory_dir, f"{task_id}.txt")
+
+    with open(memory_file_full, "w") as f:
+        f.write(memory_data)
+
+    return memory_file_full
+
+
+def load_memory[T](
+    task_id: str,
+    max_tokens: Optional[int],
+    encoder: Callable[[str], list[T]],
+    decoder: Callable[[list[T]], str],
+) -> tuple[str, str]:
+    app_dir = get_app_dir_xdg()
+    memory_dir = os.path.join(app_dir, "memory")
+    memory_file = os.path.join(memory_dir, f"{task_id}.txt")
+
+    with open(memory_file, "r") as f:
+        data = f.read()
+
+    if max_tokens:
+        toks = encoder(data)
+        if len(toks) > max_tokens:
+            toks = toks[: max(0, max_tokens - 10)]
+            data = decoder(toks)
+            data += "\n(... truncated)"
+
+    project_root_match = re.search(r"# PROJECT ROOT = \s*(.*?)\s*$", data, re.MULTILINE)
+    project_root_path = ""
+    if project_root_match:
+        matched_path = project_root_match.group(1)
+        parsed_ = shlex.split(matched_path)
+        if parsed_ and len(parsed_) == 1:
+            project_root_path = parsed_[0]
+    return project_root_path, data

wcgw/client/openai_client.py

@@ -26,6 +26,7 @@ from typer import Typer
 from ..types_ import (
     BashCommand,
     BashInteraction,
+    ContextSave,
     FileEdit,
     ReadFiles,
     ReadImage,
@@ -33,11 +34,14 @@ from ..types_ import (
     WriteIfEmpty,
 )
 from .common import CostData, History, Models, discard_input
+from .memory import load_memory
 from .openai_utils import get_input_cost, get_output_cost
 from .tools import (
     DoneFlag,
     ImageData,
+    default_enc,
     get_tool_output,
+    initialize,
     which_tool,
 )
 
@@ -117,19 +121,29 @@ def loop(
 
     history: History = []
     waiting_for_assistant = False
+
+    memory = None
     if resume:
-        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"
+        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__)
 
@@ -202,10 +216,19 @@ def loop(
             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""",
+        ),
     ]
-    uname_sysname = os.uname().sysname
-    uname_machine = os.uname().machine
 
+    initial_info = initialize(
+        os.getcwd(), [], resume if (memory and resume) else "", 8000
+    )
     system = f"""
 You're an expert software engineer with shell and code knowledge.
 
@@ -217,10 +240,7 @@ Instructions:
     - 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.
 
-System information:
-    - System: {uname_sysname}
-    - Machine: {uname_machine}
-    - Current directory: {os.getcwd()}
+{initial_info}
 
 """