codekingpro/portable-devtools

1"""File search middleware for Anthropic text editor and memory tools.2 3This module provides Glob and Grep search tools that operate on files stored4in state or filesystem.5"""6 7from __future__ import annotations8 9import fnmatch10import json11import re12import subprocess13from contextlib import suppress14from datetime import datetime, timezone15from pathlib import Path16from typing import Literal17 18from langchain_core.tools import tool19 20from langchain.agents.middleware.types import AgentMiddleware, AgentState, ContextT, ResponseT21 22 23def _expand_include_patterns(pattern: str) -> list[str] | None:24    """Expand brace patterns like `*.{py,pyi}` into a list of globs."""25    if "}" in pattern and "{" not in pattern:26        return None27 28    expanded: list[str] = []29 30    def _expand(current: str) -> None:31        start = current.find("{")32        if start == -1:33            expanded.append(current)34            return35 36        end = current.find("}", start)37        if end == -1:38            raise ValueError39 40        prefix = current[:start]41        suffix = current[end + 1 :]42        inner = current[start + 1 : end]43        if not inner:44            raise ValueError45 46        for option in inner.split(","):47            _expand(prefix + option + suffix)48 49    try:50        _expand(pattern)51    except ValueError:52        return None53 54    return expanded55 56 57def _is_valid_include_pattern(pattern: str) -> bool:58    """Validate glob pattern used for include filters."""59    if not pattern:60        return False61 62    if any(char in pattern for char in ("\x00", "\n", "\r")):63        return False64 65    expanded = _expand_include_patterns(pattern)66    if expanded is None:67        return False68 69    try:70        for candidate in expanded:71            re.compile(fnmatch.translate(candidate))72    except re.error:73        return False74 75    return True76 77 78def _match_include_pattern(basename: str, pattern: str) -> bool:79    """Return True if the basename matches the include pattern."""80    expanded = _expand_include_patterns(pattern)81    if not expanded:82        return False83 84    return any(fnmatch.fnmatch(basename, candidate) for candidate in expanded)85 86 87class FilesystemFileSearchMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]):88    """Provides Glob and Grep search over filesystem files.89 90    This middleware adds two tools that search through local filesystem:91 92    - Glob: Fast file pattern matching by file path93    - Grep: Fast content search using ripgrep or Python fallback94 95    Example:96        ```python97        from langchain.agents import create_agent98        from langchain.agents.middleware import (99            FilesystemFileSearchMiddleware,100        )101 102        agent = create_agent(103            model=model,104            tools=[],  # Add tools as needed105            middleware=[106                FilesystemFileSearchMiddleware(root_path="/workspace"),107            ],108        )109        ```110    """111 112    def __init__(113        self,114        *,115        root_path: str,116        use_ripgrep: bool = True,117        max_file_size_mb: int = 10,118    ) -> None:119        """Initialize the search middleware.120 121        Args:122            root_path: Root directory to search.123            use_ripgrep: Whether to use `ripgrep` for search.124 125                Falls back to Python if `ripgrep` unavailable.126            max_file_size_mb: Maximum file size to search in MB.127        """128        self.root_path = Path(root_path).resolve()129        self.use_ripgrep = use_ripgrep130        self.max_file_size_bytes = max_file_size_mb * 1024 * 1024131 132        # Create tool instances as closures that capture self133        @tool134        def glob_search(pattern: str, path: str = "/") -> str:135            """Fast file pattern matching tool that works with any codebase size.136 137            Supports glob patterns like `**/*.js` or `src/**/*.ts`.138 139            Returns matching file paths sorted by modification time.140 141            Use this tool when you need to find files by name patterns.142 143            Args:144                pattern: The glob pattern to match files against.145                path: The directory to search in. If not specified, searches from root.146 147            Returns:148                Newline-separated list of matching file paths, sorted by modification149                time (most recently modified first). Returns `'No files found'` if no150                matches.151            """152            try:153                base_full = self._validate_and_resolve_path(path)154            except ValueError:155                return "No files found"156 157            if not base_full.exists() or not base_full.is_dir():158                return "No files found"159 160            # Use pathlib glob161            matching: list[tuple[str, str]] = []162            for match in base_full.glob(pattern):163                if match.is_file():164                    # Convert to virtual path165                    virtual_path = "/" + str(match.relative_to(self.root_path))166                    stat = match.stat()167                    modified_at = datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc).isoformat()168                    matching.append((virtual_path, modified_at))169 170            if not matching:171                return "No files found"172 173            file_paths = [p for p, _ in matching]174            return "\n".join(file_paths)175 176        @tool177        def grep_search(178            pattern: str,179            path: str = "/",180            include: str | None = None,181            output_mode: Literal["files_with_matches", "content", "count"] = "files_with_matches",182        ) -> str:183            """Fast content search tool that works with any codebase size.184 185            Searches file contents using regular expressions. Supports full regex186            syntax and filters files by pattern with the include parameter.187 188            Args:189                pattern: The regular expression pattern to search for in file contents.190                path: The directory to search in. If not specified, searches from root.191                include: File pattern to filter (e.g., `'*.js'`, `'*.{ts,tsx}'`).192                output_mode: Output format:193 194                    - `'files_with_matches'`: Only file paths containing matches195                    - `'content'`: Matching lines with `file:line:content` format196                    - `'count'`: Count of matches per file197 198            Returns:199                Search results formatted according to `output_mode`.200                    Returns `'No matches found'` if no results.201            """202            # Compile regex pattern (for validation)203            try:204                re.compile(pattern)205            except re.error as e:206                return f"Invalid regex pattern: {e}"207 208            if include and not _is_valid_include_pattern(include):209                return "Invalid include pattern"210 211            # Try ripgrep first if enabled212            results = None213            if self.use_ripgrep:214                with suppress(215                    FileNotFoundError,216                    subprocess.CalledProcessError,217                    subprocess.TimeoutExpired,218                ):219                    results = self._ripgrep_search(pattern, path, include)220 221            # Python fallback if ripgrep failed or is disabled222            if results is None:223                results = self._python_search(pattern, path, include)224 225            if not results:226                return "No matches found"227 228            # Format output based on mode229            return self._format_grep_results(results, output_mode)230 231        self.glob_search = glob_search232        self.grep_search = grep_search233        self.tools = [glob_search, grep_search]234 235    def _validate_and_resolve_path(self, path: str) -> Path:236        """Validate and resolve a virtual path to filesystem path."""237        # Normalize path238        if not path.startswith("/"):239            path = "/" + path240 241        # Check for path traversal242        if ".." in path or "~" in path:243            msg = "Path traversal not allowed"244            raise ValueError(msg)245 246        # Convert virtual path to filesystem path247        relative = path.lstrip("/")248        full_path = (self.root_path / relative).resolve()249 250        # Ensure path is within root251        try:252            full_path.relative_to(self.root_path)253        except ValueError:254            msg = f"Path outside root directory: {path}"255            raise ValueError(msg) from None256 257        return full_path258 259    def _ripgrep_search(260        self, pattern: str, base_path: str, include: str | None261    ) -> dict[str, list[tuple[int, str]]]:262        """Search using ripgrep subprocess."""263        try:264            base_full = self._validate_and_resolve_path(base_path)265        except ValueError:266            return {}267 268        if not base_full.exists():269            return {}270 271        # Build ripgrep command272        cmd = ["rg", "--json"]273 274        if include:275            # Convert glob pattern to ripgrep glob276            cmd.extend(["--glob", include])277 278        cmd.extend(["--", pattern, str(base_full)])279 280        try:281            result = subprocess.run(  # noqa: S603282                cmd,283                capture_output=True,284                text=True,285                timeout=30,286                check=False,287            )288        except (subprocess.TimeoutExpired, FileNotFoundError):289            # Fallback to Python search if ripgrep unavailable or times out290            return self._python_search(pattern, base_path, include)291 292        # Parse ripgrep JSON output293        results: dict[str, list[tuple[int, str]]] = {}294        for line in result.stdout.splitlines():295            try:296                data = json.loads(line)297                if data["type"] == "match":298                    path = data["data"]["path"]["text"]299                    # Convert to virtual path300                    virtual_path = "/" + str(Path(path).relative_to(self.root_path))301                    line_num = data["data"]["line_number"]302                    line_text = data["data"]["lines"]["text"].rstrip("\n")303 304                    if virtual_path not in results:305                        results[virtual_path] = []306                    results[virtual_path].append((line_num, line_text))307            except (json.JSONDecodeError, KeyError):308                continue309 310        return results311 312    def _python_search(313        self, pattern: str, base_path: str, include: str | None314    ) -> dict[str, list[tuple[int, str]]]:315        """Search using Python regex (fallback)."""316        try:317            base_full = self._validate_and_resolve_path(base_path)318        except ValueError:319            return {}320 321        if not base_full.exists():322            return {}323 324        regex = re.compile(pattern)325        results: dict[str, list[tuple[int, str]]] = {}326 327        # Walk directory tree328        for file_path in base_full.rglob("*"):329            if not file_path.is_file():330                continue331 332            # Check include filter333            if include and not _match_include_pattern(file_path.name, include):334                continue335 336            # Skip files that are too large337            if file_path.stat().st_size > self.max_file_size_bytes:338                continue339 340            try:341                content = file_path.read_text()342            except (UnicodeDecodeError, PermissionError):343                continue344 345            # Search content346            for line_num, line in enumerate(content.splitlines(), 1):347                if regex.search(line):348                    virtual_path = "/" + str(file_path.relative_to(self.root_path))349                    if virtual_path not in results:350                        results[virtual_path] = []351                    results[virtual_path].append((line_num, line))352 353        return results354 355    @staticmethod356    def _format_grep_results(357        results: dict[str, list[tuple[int, str]]],358        output_mode: str,359    ) -> str:360        """Format grep results based on output mode."""361        if output_mode == "files_with_matches":362            # Just return file paths363            return "\n".join(sorted(results.keys()))364 365        if output_mode == "content":366            # Return file:line:content format367            lines = []368            for file_path in sorted(results.keys()):369                for line_num, line in results[file_path]:370                    lines.append(f"{file_path}:{line_num}:{line}")371            return "\n".join(lines)372 373        if output_mode == "count":374            # Return file:count format375            lines = []376            for file_path in sorted(results.keys()):377                count = len(results[file_path])378                lines.append(f"{file_path}:{count}")379            return "\n".join(lines)380 381        # Default to files_with_matches382        return "\n".join(sorted(results.keys()))383 384 385__all__ = [386    "FilesystemFileSearchMiddleware",387]388