hanzo-mcp 0.5.2__py3-none-any.whl → 0.6.1__py3-none-any.whl

hanzo_mcp/tools/filesystem/unified_search.py

@@ -0,0 +1,711 @@
+"""Unified search tool that runs multiple search types in parallel.
+
+This tool consolidates all search capabilities and runs them concurrently:
+- grep: Fast pattern/regex search using ripgrep
+- grep_ast: AST-aware code search with structural context
+- vector_search: Semantic similarity search
+- git_search: Search through git history
+- symbol_search: Find symbols (functions, classes) in code
+
+Results are combined, deduplicated, and ranked for comprehensive search coverage.
+"""
+
+import asyncio
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Annotated, Dict, List, Optional, Set, Tuple, TypedDict, Unpack, final, override
+from enum import Enum
+
+from mcp.server.fastmcp import Context as MCPContext
+from mcp.server import FastMCP
+from pydantic import Field
+
+from hanzo_mcp.tools.filesystem.base import FilesystemBaseTool
+from hanzo_mcp.tools.filesystem.grep import Grep
+from hanzo_mcp.tools.filesystem.symbols import SymbolsTool
+from hanzo_mcp.tools.filesystem.git_search import GitSearchTool
+from hanzo_mcp.tools.vector.vector_search import VectorSearchTool
+from hanzo_mcp.tools.common.permissions import PermissionManager
+from hanzo_mcp.tools.vector.project_manager import ProjectVectorManager
+
+
+class SearchType(Enum):
+    """Types of searches that can be performed."""
+    GREP = "grep"
+    GREP_AST = "grep_ast"
+    VECTOR = "vector"
+    GIT = "git"
+    SYMBOL = "symbol"  # Searches for function/class definitions
+
+
+@dataclass
+class SearchResult:
+    """Unified search result from any search type."""
+    file_path: str
+    line_number: Optional[int]
+    content: str
+    search_type: SearchType
+    score: float  # Relevance score (0-1)
+    context: Optional[str] = None  # Function/class context
+    match_count: int = 1  # Number of matches in this location
+
+
+Pattern = Annotated[
+    str,
+    Field(
+        description="The search pattern (supports regex for grep, natural language for vector search)",
+        min_length=1,
+    ),
+]
+
+SearchPath = Annotated[
+    str,
+    Field(
+        description="The directory to search in. Defaults to current directory.",
+        default=".",
+    ),
+]
+
+Include = Annotated[
+    str,
+    Field(
+        description='File pattern to include (e.g. "*.js", "*.{ts,tsx}")',
+        default="*",
+    ),
+]
+
+MaxResults = Annotated[
+    int,
+    Field(
+        description="Maximum number of results to return",
+        default=50,
+    ),
+]
+
+EnableGrep = Annotated[
+    bool,
+    Field(
+        description="Enable fast pattern/regex search",
+        default=True,
+    ),
+]
+
+EnableGrepAst = Annotated[
+    bool,
+    Field(
+        description="Enable AST-aware search with code structure context",
+        default=True,
+    ),
+]
+
+EnableVector = Annotated[
+    bool,
+    Field(
+        description="Enable semantic similarity search",
+        default=True,
+    ),
+]
+
+EnableGit = Annotated[
+    bool,
+    Field(
+        description="Enable git history search",
+        default=True,
+    ),
+]
+
+EnableSymbol = Annotated[
+    bool,
+    Field(
+        description="Enable symbol search (functions, classes)",
+        default=True,
+    ),
+]
+
+IncludeContext = Annotated[
+    bool,
+    Field(
+        description="Include function/class context for matches",
+        default=True,
+    ),
+]
+
+
+class UnifiedSearchParams(TypedDict):
+    """Parameters for unified search."""
+    pattern: Pattern
+    path: SearchPath
+    include: Include
+    max_results: MaxResults
+    enable_grep: EnableGrep
+    enable_grep_ast: EnableGrepAst
+    enable_vector: EnableVector
+    enable_git: EnableGit
+    enable_symbol: EnableSymbol
+    include_context: IncludeContext
+
+
+@final
+class UnifiedSearchTool(FilesystemBaseTool):
+    """Unified search tool that runs multiple search types in parallel."""
+    
+    def __init__(self, permission_manager: PermissionManager, 
+                 project_manager: Optional[ProjectVectorManager] = None):
+        """Initialize the unified search tool.
+        
+        Args:
+            permission_manager: Permission manager for access control
+            project_manager: Optional project manager for vector search
+        """
+        super().__init__(permission_manager)
+        self.project_manager = project_manager
+        
+        # Initialize component tools
+        self.grep_tool = Grep(permission_manager)
+        self.grep_ast_tool = SymbolsTool(permission_manager)
+        self.git_search_tool = GitSearchTool(permission_manager)
+        
+        # Vector search is optional
+        self.vector_tool = None
+        if project_manager:
+            self.vector_tool = VectorSearchTool(permission_manager, project_manager)
+    
+    @property
+    @override
+    def name(self) -> str:
+        """Get the tool name."""
+        return "unified_search"
+    
+    @property
+    @override
+    def description(self) -> str:
+        """Get the tool description."""
+        return """Unified search that runs multiple search strategies in parallel.
+
+Automatically runs the most appropriate search types based on your pattern:
+- Pattern matching (grep) for exact text/regex
+- AST search for code structure understanding  
+- Semantic search for concepts and meaning
+- Git history for tracking changes
+- Symbol search for finding definitions
+
+All searches run concurrently for maximum speed. Results are combined,
+deduplicated, and ranked by relevance.
+
+Examples:
+- Search for TODO comments: pattern="TODO"
+- Find error handling: pattern="error handling implementation"
+- Locate function: pattern="processPayment"
+- Track changes: pattern="bug fix" (searches git history too)
+
+This is the recommended search tool for comprehensive results."""
+    
+    def _analyze_pattern(self, pattern: str) -> Dict[str, bool]:
+        """Analyze the pattern to determine optimal search strategies.
+        
+        Args:
+            pattern: The search pattern
+            
+        Returns:
+            Dictionary of search type recommendations
+        """
+        # Check if pattern looks like regex
+        regex_chars = r'[.*+?^${}()|[\]\\]'
+        has_regex = bool(re.search(regex_chars, pattern))
+        
+        # Check if pattern looks like a symbol name
+        is_symbol = bool(re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', pattern))
+        
+        # Check if pattern is natural language
+        words = pattern.split()
+        is_natural_language = len(words) > 2 and not has_regex
+        
+        return {
+            'use_grep': True,  # Always useful
+            'use_grep_ast': not has_regex,  # AST doesn't handle regex well
+            'use_vector': is_natural_language or len(pattern) > 10,
+            'use_git': True,  # Always check history
+            'use_symbol': is_symbol or 'def' in pattern or 'class' in pattern
+        }
+    
+    async def _run_grep_search(self, pattern: str, path: str, include: str,
+                              tool_ctx, max_results: int) -> List[SearchResult]:
+        """Run grep search and parse results."""
+        try:
+            result = await self.grep_tool.call(
+                tool_ctx.mcp_context,
+                pattern=pattern,
+                path=path,
+                include=include
+            )
+            
+            results = []
+            if "Found" in result and "matches" in result:
+                lines = result.split('\n')
+                for line in lines[2:]:  # Skip header
+                    if ':' in line and line.strip():
+                        try:
+                            parts = line.split(':', 2)
+                            if len(parts) >= 3:
+                                results.append(SearchResult(
+                                    file_path=parts[0],
+                                    line_number=int(parts[1]),
+                                    content=parts[2].strip(),
+                                    search_type=SearchType.GREP,
+                                    score=1.0  # Exact matches get perfect score
+                                ))
+                                if len(results) >= max_results:
+                                    break
+                        except ValueError:
+                            continue
+            
+            await tool_ctx.info(f"Grep found {len(results)} results")
+            return results
+            
+        except Exception as e:
+            await tool_ctx.error(f"Grep search failed: {e}")
+            return []
+    
+    async def _run_grep_ast_search(self, pattern: str, path: str,
+                                  tool_ctx, max_results: int) -> List[SearchResult]:
+        """Run AST-aware search and parse results."""
+        try:
+            result = await self.grep_ast_tool.call(
+                tool_ctx.mcp_context,
+                pattern=pattern,
+                path=path,
+                ignore_case=True,
+                line_number=True
+            )
+            
+            results = []
+            if result and not result.startswith("No matches"):
+                current_file = None
+                current_context = []
+                
+                for line in result.split('\n'):
+                    if line.endswith(':') and '/' in line:
+                        current_file = line[:-1]
+                        current_context = []
+                    elif current_file and ':' in line:
+                        try:
+                            # Try to parse line with number
+                            parts = line.split(':', 1)
+                            line_num = int(parts[0].strip())
+                            content = parts[1].strip() if len(parts) > 1 else ""
+                            
+                            results.append(SearchResult(
+                                file_path=current_file,
+                                line_number=line_num,
+                                content=content,
+                                search_type=SearchType.GREP_AST,
+                                score=0.95,  # High score for AST matches
+                                context=" > ".join(current_context) if current_context else None
+                            ))
+                            
+                            if len(results) >= max_results:
+                                break
+                        except ValueError:
+                            # This might be context info
+                            if line.strip():
+                                current_context.append(line.strip())
+            
+            await tool_ctx.info(f"AST search found {len(results)} results")
+            return results
+            
+        except Exception as e:
+            await tool_ctx.error(f"AST search failed: {e}")
+            return []
+    
+    async def _run_vector_search(self, pattern: str, path: str,
+                                tool_ctx, max_results: int) -> List[SearchResult]:
+        """Run semantic vector search."""
+        if not self.vector_tool:
+            return []
+            
+        try:
+            # Determine search scope
+            search_scope = "current" if path == "." else "all"
+            
+            result = await self.vector_tool.call(
+                tool_ctx.mcp_context,
+                query=pattern,
+                limit=max_results,
+                score_threshold=0.3,
+                search_scope=search_scope,
+                include_content=True
+            )
+            
+            results = []
+            if "Found" in result:
+                # Parse vector search results
+                lines = result.split('\n')
+                current_file = None
+                current_score = 0.0
+                
+                for line in lines:
+                    if "Result" in line and "Score:" in line:
+                        # Extract score and file
+                        score_match = re.search(r'Score: ([\d.]+)%', line)
+                        if score_match:
+                            current_score = float(score_match.group(1)) / 100.0
+                        
+                        file_match = re.search(r' - ([^\s]+)$', line)
+                        if file_match:
+                            current_file = file_match.group(1)
+                    
+                    elif current_file and line.strip() and not line.startswith('-'):
+                        # Content line
+                        results.append(SearchResult(
+                            file_path=current_file,
+                            line_number=None,
+                            content=line.strip()[:200],  # Limit content length
+                            search_type=SearchType.VECTOR,
+                            score=current_score
+                        ))
+                        
+                        if len(results) >= max_results:
+                            break
+            
+            await tool_ctx.info(f"Vector search found {len(results)} results")
+            return results
+            
+        except Exception as e:
+            await tool_ctx.error(f"Vector search failed: {e}")
+            return []
+    
+    async def _run_git_search(self, pattern: str, path: str,
+                             tool_ctx, max_results: int) -> List[SearchResult]:
+        """Run git history search."""
+        try:
+            # Search in both content and commits
+            tasks = [
+                self.git_search_tool.call(
+                    tool_ctx.mcp_context,
+                    pattern=pattern,
+                    path=path,
+                    search_type="content",
+                    max_count=max_results // 2
+                ),
+                self.git_search_tool.call(
+                    tool_ctx.mcp_context,
+                    pattern=pattern,
+                    path=path,
+                    search_type="commits",
+                    max_count=max_results // 2
+                )
+            ]
+            
+            git_results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            results = []
+            for i, result in enumerate(git_results):
+                if isinstance(result, Exception):
+                    continue
+                    
+                if "Found" in result:
+                    # Parse git results
+                    lines = result.split('\n')
+                    for line in lines:
+                        if ':' in line and line.strip():
+                            parts = line.split(':', 2)
+                            if len(parts) >= 2:
+                                results.append(SearchResult(
+                                    file_path=parts[0].strip(),
+                                    line_number=None,
+                                    content=parts[-1].strip() if len(parts) > 2 else line,
+                                    search_type=SearchType.GIT,
+                                    score=0.8  # Good score for git matches
+                                ))
+                                
+                                if len(results) >= max_results:
+                                    break
+            
+            await tool_ctx.info(f"Git search found {len(results)} results")
+            return results
+            
+        except Exception as e:
+            await tool_ctx.error(f"Git search failed: {e}")
+            return []
+    
+    async def _run_symbol_search(self, pattern: str, path: str,
+                                tool_ctx, max_results: int) -> List[SearchResult]:
+        """Search for symbol definitions using grep with specific patterns."""
+        try:
+            # Create patterns for common symbol definitions
+            symbol_patterns = [
+                f"(def|class|function|func|fn)\\s+{pattern}",  # Python, JS, various
+                f"(public|private|protected)?\\s*(static)?\\s*\\w+\\s+{pattern}\\s*\\(",  # Java/C++
+                f"const\\s+{pattern}\\s*=",  # JS/TS const
+                f"let\\s+{pattern}\\s*=",  # JS/TS let
+                f"var\\s+{pattern}\\s*=",  # JS/TS var
+            ]
+            
+            # Run grep searches in parallel for each pattern
+            tasks = []
+            for sp in symbol_patterns:
+                tasks.append(
+                    self.grep_tool.call(
+                        tool_ctx.mcp_context,
+                        pattern=sp,
+                        path=path,
+                        include="*"
+                    )
+                )
+            
+            grep_results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            results = []
+            for result in grep_results:
+                if isinstance(result, Exception):
+                    continue
+                    
+                if "Found" in result and "matches" in result:
+                    lines = result.split('\n')
+                    for line in lines[2:]:  # Skip header
+                        if ':' in line and line.strip():
+                            try:
+                                parts = line.split(':', 2)
+                                if len(parts) >= 3:
+                                    results.append(SearchResult(
+                                        file_path=parts[0],
+                                        line_number=int(parts[1]),
+                                        content=parts[2].strip(),
+                                        search_type=SearchType.SYMBOL,
+                                        score=0.98  # Very high score for symbol definitions
+                                    ))
+                                    if len(results) >= max_results:
+                                        break
+                            except ValueError:
+                                continue
+            
+            await tool_ctx.info(f"Symbol search found {len(results)} results")
+            return results
+            
+        except Exception as e:
+            await tool_ctx.error(f"Symbol search failed: {e}")
+            return []
+    
+    def _deduplicate_results(self, all_results: List[SearchResult]) -> List[SearchResult]:
+        """Deduplicate results, keeping the highest scoring version."""
+        seen = {}
+        
+        for result in all_results:
+            key = (result.file_path, result.line_number)
+            
+            if key not in seen or result.score > seen[key].score:
+                seen[key] = result
+            elif key in seen and result.context and not seen[key].context:
+                # Add context if missing
+                seen[key].context = result.context
+        
+        return list(seen.values())
+    
+    def _rank_results(self, results: List[SearchResult]) -> List[SearchResult]:
+        """Rank results by relevance score and search type priority."""
+        # Define search type priorities
+        type_priority = {
+            SearchType.SYMBOL: 5,
+            SearchType.GREP: 4,
+            SearchType.GREP_AST: 3,
+            SearchType.GIT: 2,
+            SearchType.VECTOR: 1
+        }
+        
+        # Sort by score (descending) and then by type priority
+        results.sort(
+            key=lambda r: (r.score, type_priority.get(r.search_type, 0)),
+            reverse=True
+        )
+        
+        return results
+    
+    @override
+    async def call(
+        self,
+        ctx: MCPContext,
+        **params: Unpack[UnifiedSearchParams],
+    ) -> str:
+        """Execute unified search across all enabled search types."""
+        import time
+        start_time = time.time()
+        
+        tool_ctx = self.create_tool_context(ctx)
+        
+        # Extract parameters
+        pattern = params["pattern"]
+        path = params.get("path", ".")
+        include = params.get("include", "*")
+        max_results = params.get("max_results", 50)
+        include_context = params.get("include_context", True)
+        
+        # Validate path
+        path_validation = self.validate_path(path)
+        if path_validation.is_error:
+            await tool_ctx.error(path_validation.error_message)
+            return f"Error: {path_validation.error_message}"
+        
+        # Check permissions
+        allowed, error_msg = await self.check_path_allowed(path, tool_ctx)
+        if not allowed:
+            return error_msg
+        
+        # Check existence
+        exists, error_msg = await self.check_path_exists(path, tool_ctx)
+        if not exists:
+            return error_msg
+        
+        # Analyze pattern to determine best search strategies
+        pattern_analysis = self._analyze_pattern(pattern)
+        
+        await tool_ctx.info(f"Starting unified search for '{pattern}' in {path}")
+        
+        # Build list of search tasks based on enabled types and pattern analysis
+        search_tasks = []
+        search_names = []
+        
+        if params.get("enable_grep", True) and pattern_analysis['use_grep']:
+            search_tasks.append(self._run_grep_search(pattern, path, include, tool_ctx, max_results))
+            search_names.append("grep")
+        
+        if params.get("enable_grep_ast", True) and pattern_analysis['use_grep_ast']:
+            search_tasks.append(self._run_grep_ast_search(pattern, path, tool_ctx, max_results))
+            search_names.append("grep_ast")
+        
+        if params.get("enable_vector", True) and self.vector_tool and pattern_analysis['use_vector']:
+            search_tasks.append(self._run_vector_search(pattern, path, tool_ctx, max_results))
+            search_names.append("vector")
+        
+        if params.get("enable_git", True) and pattern_analysis['use_git']:
+            search_tasks.append(self._run_git_search(pattern, path, tool_ctx, max_results))
+            search_names.append("git")
+        
+        if params.get("enable_symbol", True) and pattern_analysis['use_symbol']:
+            search_tasks.append(self._run_symbol_search(pattern, path, tool_ctx, max_results))
+            search_names.append("symbol")
+        
+        await tool_ctx.info(f"Running {len(search_tasks)} search types in parallel: {', '.join(search_names)}")
+        
+        # Run all searches in parallel
+        search_results = await asyncio.gather(*search_tasks, return_exceptions=True)
+        
+        # Collect all results
+        all_results = []
+        results_by_type = {}
+        
+        for search_type, results in zip(search_names, search_results):
+            if isinstance(results, Exception):
+                await tool_ctx.error(f"{search_type} search failed: {results}")
+                results_by_type[search_type] = []
+            else:
+                results_by_type[search_type] = results
+                all_results.extend(results)
+        
+        # Deduplicate and rank results
+        unique_results = self._deduplicate_results(all_results)
+        ranked_results = self._rank_results(unique_results)
+        
+        # Limit total results
+        final_results = ranked_results[:max_results]
+        
+        # Calculate search time
+        search_time = (time.time() - start_time) * 1000
+        
+        # Format output
+        return self._format_results(
+            pattern=pattern,
+            results=final_results,
+            results_by_type=results_by_type,
+            search_time_ms=search_time,
+            include_context=include_context
+        )
+    
+    def _format_results(self, pattern: str, results: List[SearchResult],
+                       results_by_type: Dict[str, List[SearchResult]],
+                       search_time_ms: float, include_context: bool) -> str:
+        """Format search results for display."""
+        output = []
+        
+        # Header
+        output.append(f"=== Unified Search Results ===")
+        output.append(f"Pattern: '{pattern}'")
+        output.append(f"Total results: {len(results)}")
+        output.append(f"Search time: {search_time_ms:.1f}ms")
+        
+        # Summary by type
+        output.append("\nResults by type:")
+        for search_type, type_results in results_by_type.items():
+            if type_results:
+                output.append(f"  {search_type}: {len(type_results)} matches")
+        
+        if not results:
+            output.append("\nNo results found.")
+            return "\n".join(output)
+        
+        # Group results by file
+        results_by_file = {}
+        for result in results:
+            if result.file_path not in results_by_file:
+                results_by_file[result.file_path] = []
+            results_by_file[result.file_path].append(result)
+        
+        # Display results
+        output.append(f"\n=== Results ({len(results)} total) ===\n")
+        
+        for file_path, file_results in results_by_file.items():
+            output.append(f"{file_path}")
+            output.append("-" * len(file_path))
+            
+            # Sort by line number
+            file_results.sort(key=lambda r: r.line_number or 0)
+            
+            for result in file_results:
+                # Format result line
+                score_str = f"[{result.search_type.value} {result.score:.2f}]"
+                
+                if result.line_number:
+                    output.append(f"  {result.line_number:>4}: {score_str} {result.content}")
+                else:
+                    output.append(f"       {score_str} {result.content}")
+                
+                # Add context if available and requested
+                if include_context and result.context:
+                    output.append(f"         Context: {result.context}")
+            
+            output.append("")  # Empty line between files
+        
+        return "\n".join(output)
+    
+    @override
+    def register(self, mcp_server: FastMCP) -> None:
+        """Register the unified search tool with the MCP server."""
+        tool_self = self
+        
+        @mcp_server.tool(name=self.name, description=self.description)
+        async def unified_search(
+            ctx: MCPContext,
+            pattern: Pattern,
+            path: SearchPath = ".",
+            include: Include = "*",
+            max_results: MaxResults = 50,
+            enable_grep: EnableGrep = True,
+            enable_grep_ast: EnableGrepAst = True,
+            enable_vector: EnableVector = True,
+            enable_git: EnableGit = True,
+            enable_symbol: EnableSymbol = True,
+            include_context: IncludeContext = True,
+        ) -> str:
+            return await tool_self.call(
+                ctx,
+                pattern=pattern,
+                path=path,
+                include=include,
+                max_results=max_results,
+                enable_grep=enable_grep,
+                enable_grep_ast=enable_grep_ast,
+                enable_vector=enable_vector,
+                enable_git=enable_git,
+                enable_symbol=enable_symbol,
+                include_context=include_context,
+            )