Add global replacement tool

src/codegen/extensions/langchain/agent.py

@@ -1,6 +1,6 @@
 """Demo implementation of an agent with Codegen tools."""
 
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 from langchain.tools import BaseTool
 from langchain_core.messages import SystemMessage
@@ -13,13 +13,15 @@
 from codegen.extensions.langchain.tools import (
     CreateFileTool,
     DeleteFileTool,
+    GlobalReplacementEditTool,
     ListDirectoryTool,
     MoveSymbolTool,
     ReflectionTool,
     RelaceEditTool,
     RenameFileTool,
     ReplacementEditTool,
     RevealSymbolTool,
+    SearchFilesByNameTool,
     SearchTool,
     # SemanticEditTool,
     ViewFileTool,
@@ -38,8 +40,8 @@ def create_codebase_agent(
     system_message: SystemMessage = SystemMessage(REASONER_SYSTEM_MESSAGE),
     memory: bool = True,
     debug: bool = False,
-    additional_tools: Optional[list[BaseTool]] = None,
-    config: Optional[AgentConfig] = None,
+    additional_tools: list[BaseTool] | None = None,
+    config: AgentConfig | None = None,
     **kwargs,
 ) -> CompiledGraph:
     """Create an agent with all codebase tools.
@@ -76,6 +78,8 @@ def create_codebase_agent(
         ReplacementEditTool(codebase),
         RelaceEditTool(codebase),
         ReflectionTool(codebase),
+        SearchFilesByNameTool(codebase),
+        GlobalReplacementEditTool(codebase),
         # SemanticSearchTool(codebase),
         # =====[ Github Integration ]=====
         # Enable Github integration
@@ -101,8 +105,8 @@ def create_chat_agent(
     system_message: SystemMessage = SystemMessage(REASONER_SYSTEM_MESSAGE),
     memory: bool = True,
     debug: bool = False,
-    additional_tools: Optional[list[BaseTool]] = None,
-    config: Optional[dict[str, Any]] = None,  # over here you can pass in the max length of the number of messages
+    additional_tools: list[BaseTool] | None = None,
+    config: dict[str, Any] | None = None,  # over here you can pass in the max length of the number of messages
     **kwargs,
 ) -> CompiledGraph:
     """Create an agent with all codebase tools.
@@ -151,7 +155,7 @@ def create_codebase_inspector_agent(
     system_message: SystemMessage = SystemMessage(REASONER_SYSTEM_MESSAGE),
     memory: bool = True,
     debug: bool = True,
-    config: Optional[dict[str, Any]] = None,
+    config: dict[str, Any] | None = None,
     **kwargs,
 ) -> CompiledGraph:
     """Create an inspector agent with read-only codebase tools.
@@ -189,7 +193,7 @@ def create_agent_with_tools(
     system_message: SystemMessage = SystemMessage(REASONER_SYSTEM_MESSAGE),
     memory: bool = True,
     debug: bool = True,
-    config: Optional[dict[str, Any]] = None,
+    config: dict[str, Any] | None = None,
     **kwargs,
 ) -> CompiledGraph:
     """Create an agent with a specific set of tools.

src/codegen/extensions/langchain/tools.py

@@ -1,6 +1,7 @@
 """Langchain tools for workspace operations."""
 
-from typing import Callable, ClassVar, Literal, Optional
+from collections.abc import Callable
+from typing import ClassVar, Literal
 
 from langchain_core.tools.base import BaseTool
 from pydantic import BaseModel, Field
@@ -9,6 +10,7 @@
 from codegen.extensions.tools.bash import run_bash_command
 from codegen.extensions.tools.github.checkout_pr import checkout_pr
 from codegen.extensions.tools.github.view_pr_checks import view_pr_checks
+from codegen.extensions.tools.global_replacement_edit import replacement_edit_global
 from codegen.extensions.tools.linear.linear import (
     linear_comment_on_issue_tool,
     linear_create_issue_tool,
@@ -50,10 +52,10 @@ class ViewFileInput(BaseModel):
     """Input for viewing a file."""
 
     filepath: str = Field(..., description="Path to the file relative to workspace root")
-    start_line: Optional[int] = Field(None, description="Starting line number to view (1-indexed, inclusive)")
-    end_line: Optional[int] = Field(None, description="Ending line number to view (1-indexed, inclusive)")
-    max_lines: Optional[int] = Field(None, description="Maximum number of lines to view at once, defaults to 250")
-    line_numbers: Optional[bool] = Field(True, description="If True, add line numbers to the content (1-indexed)")
+    start_line: int | None = Field(None, description="Starting line number to view (1-indexed, inclusive)")
+    end_line: int | None = Field(None, description="Ending line number to view (1-indexed, inclusive)")
+    max_lines: int | None = Field(None, description="Maximum number of lines to view at once, defaults to 250")
+    line_numbers: bool | None = Field(True, description="If True, add line numbers to the content (1-indexed)")
 
 
 class ViewFileTool(BaseTool):
@@ -72,10 +74,10 @@ def __init__(self, codebase: Codebase) -> None:
     def _run(
         self,
         filepath: str,
-        start_line: Optional[int] = None,
-        end_line: Optional[int] = None,
-        max_lines: Optional[int] = None,
-        line_numbers: Optional[bool] = True,
+        start_line: int | None = None,
+        end_line: int | None = None,
+        max_lines: int | None = None,
+        line_numbers: bool | None = True,
     ) -> str:
         result = view_file(
             self.codebase,
@@ -120,7 +122,7 @@ class SearchInput(BaseModel):
         description="""The search query to find in the codebase. When ripgrep is available, this will be passed as a ripgrep pattern. For regex searches, set use_regex=True.
         Ripgrep is the preferred method.""",
     )
-    file_extensions: Optional[list[str]] = Field(default=None, description="Optional list of file extensions to search (e.g. ['.py', '.ts'])")
+    file_extensions: list[str] | None = Field(default=None, description="Optional list of file extensions to search (e.g. ['.py', '.ts'])")
     page: int = Field(default=1, description="Page number to return (1-based, default: 1)")
     files_per_page: int = Field(default=10, description="Number of files to return per page (default: 10)")
     use_regex: bool = Field(default=False, description="Whether to treat query as a regex pattern (default: False)")
@@ -137,7 +139,7 @@ class SearchTool(BaseTool):
     def __init__(self, codebase: Codebase) -> None:
         super().__init__(codebase=codebase)
 
-    def _run(self, query: str, file_extensions: Optional[list[str]] = None, page: int = 1, files_per_page: int = 10, use_regex: bool = False) -> str:
+    def _run(self, query: str, file_extensions: list[str] | None = None, page: int = 1, files_per_page: int = 10, use_regex: bool = False) -> str:
         result = search(self.codebase, query, file_extensions=file_extensions, page=page, files_per_page=files_per_page, use_regex=use_regex)
         return result.render()
 
@@ -273,7 +275,7 @@ class RevealSymbolInput(BaseModel):
 
     symbol_name: str = Field(..., description="Name of the symbol to analyze")
     degree: int = Field(default=1, description="How many degrees of separation to traverse")
-    max_tokens: Optional[int] = Field(
+    max_tokens: int | None = Field(
         default=None,
         description="Optional maximum number of tokens for all source code combined",
     )
@@ -296,7 +298,7 @@ def _run(
         self,
         symbol_name: str,
         degree: int = 1,
-        max_tokens: Optional[int] = None,
+        max_tokens: int | None = None,
         collect_dependencies: bool = True,
         collect_usages: bool = True,
     ) -> str:
@@ -849,8 +851,10 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]:
         RenameFileTool(codebase),
         ReplacementEditTool(codebase),
         RevealSymbolTool(codebase),
+        GlobalReplacementEditTool(codebase),
         RunBashCommandTool(),  # Note: This tool doesn't need the codebase
         SearchTool(codebase),
+        SearchFilesByNameTool(codebase),
         # SemanticEditTool(codebase),
         # SemanticSearchTool(codebase),
         ViewFileTool(codebase),
@@ -872,6 +876,62 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]:
     ]
 
 
+class GlobalReplacementEditInput(BaseModel):
+    """Input for replacement editing across the entire codebase."""
+
+    file_pattern: str = Field(
+        default="*",
+        description=("Glob pattern to match files that should be edited. Supports all Python glob syntax including wildcards (*, ?, **)"),
+    )
+    pattern: str = Field(
+        ...,
+        description=(
+            "Regular expression pattern to match text that should be replaced. "
+            "Supports all Python regex syntax including capture groups (\\1, \\2, etc). "
+            "The pattern is compiled with re.MULTILINE flag by default."
+        ),
+    )
+    replacement: str = Field(
+        ...,
+        description=(
+            "Text to replace matched patterns with. Can reference regex capture groups using \\1, \\2, etc. If using regex groups in pattern, make sure to preserve them in replacement if needed."
+        ),
+    )
+    count: int | None = Field(
+        default=None,
+        description=(
+            "Maximum number of replacements to make. "
+            "Use None to replace all occurrences (default), or specify a number to limit replacements. "
+            "Useful when you only want to replace the first N occurrences."
+        ),
+    )
+
+
+class GlobalReplacementEditTool(BaseTool):
+    """Tool for regex-based replacement editing of files across the entire codebase.
+
+    Use this to make a change across an entire codebase if you have a regex pattern that matches the text you want to replace and are trying to edit a large number of files.
+    """
+
+    name: ClassVar[str] = "global_replace"
+    description: ClassVar[str] = "Replace text in the entire codebase using regex pattern matching."
+    args_schema: ClassVar[type[BaseModel]] = GlobalReplacementEditInput
+    codebase: Codebase = Field(exclude=True)
+
+    def __init__(self, codebase: Codebase) -> None:
+        super().__init__(codebase=codebase)
+
+    def _run(
+        self,
+        file_pattern: str,
+        pattern: str,
+        replacement: str,
+        count: int | None = None,
+    ) -> str:
+        result = replacement_edit_global(self.codebase, file_pattern, pattern, replacement, count)
+        return result.render()
+
+
 class ReplacementEditInput(BaseModel):
     """Input for replacement editing."""
 
@@ -905,7 +965,7 @@ class ReplacementEditInput(BaseModel):
             "Default is -1 (end of file)."
         ),
     )
-    count: Optional[int] = Field(
+    count: int | None = Field(
         default=None,
         description=(
             "Maximum number of replacements to make. "
@@ -933,7 +993,7 @@ def _run(
         replacement: str,
         start: int = 1,
         end: int = -1,
-        count: Optional[int] = None,
+        count: int | None = None,
     ) -> str:
         result = replacement_edit(
             self.codebase,
@@ -997,7 +1057,7 @@ class ReflectionInput(BaseModel):
     context_summary: str = Field(..., description="Summary of the current context and problem being solved")
     findings_so_far: str = Field(..., description="Key information and insights gathered so far")
     current_challenges: str = Field(default="", description="Current obstacles or questions that need to be addressed")
-    reflection_focus: Optional[str] = Field(default=None, description="Optional specific aspect to focus reflection on (e.g., 'architecture', 'performance', 'next steps')")
+    reflection_focus: str | None = Field(default=None, description="Optional specific aspect to focus reflection on (e.g., 'architecture', 'performance', 'next steps')")
 
 
 class ReflectionTool(BaseTool):
@@ -1020,7 +1080,7 @@ def _run(
         context_summary: str,
         findings_so_far: str,
         current_challenges: str = "",
-        reflection_focus: Optional[str] = None,
+        reflection_focus: str | None = None,
     ) -> str:
         result = perform_reflection(context_summary=context_summary, findings_so_far=findings_so_far, current_challenges=current_challenges, reflection_focus=reflection_focus, codebase=self.codebase)
 

src/codegen/extensions/tools/__init__.py

@@ -8,6 +8,7 @@
 from .github.create_pr_comment import create_pr_comment
 from .github.create_pr_review_comment import create_pr_review_comment
 from .github.view_pr import view_pr
+from .global_replacement_edit import replacement_edit_global
 from .linear import (
     linear_comment_on_issue_tool,
     linear_get_issue_comments_tool,
@@ -49,6 +50,7 @@
     "perform_reflection",
     "rename_file",
     "replacement_edit",
+    "replacement_edit_global",
     "reveal_symbol",
     "run_codemod",
     # Search operations