code-finder 0.1.0__py3-none-any.whl

claude_context/search.py

@@ -0,0 +1,493 @@
+"""
+Hybrid Search for Claude Context
+
+Combines BM25 (keyword) and vector (semantic) search for optimal code retrieval.
+This gives us the best of both worlds:
+- BM25: Exact matches, variable names, specific terms
+- Vector: Conceptual similarity, related functionality
+"""
+
+import logging
+import time
+from typing import List, Dict, Any, Optional, Tuple
+from dataclasses import dataclass, field
+from pathlib import Path
+import re
+
+import numpy as np
+from rank_bm25 import BM25Okapi
+
+from .config import ClaudeContextConfig, MilvusManager
+from .embeddings import LocalEmbeddings
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SearchResult:
+    """Represents a search result with combined scoring"""
+    content: str
+    file_path: str
+    file_name: str
+    start_line: int
+    end_line: int
+    language: str
+    chunk_type: str  # function, class, method, import_block, etc.
+    chunk_name: Optional[str] = None
+    parent_context: Optional[str] = None  # Full scope chain (e.g., "ClassName > method_name")
+
+    # Enhanced metadata (from EnhancedASTChunker)
+    signature: Optional[str] = None  # Full function signature with params and types
+    docstring: Optional[str] = None  # Extracted docstring
+    imports: Optional[str] = None  # JSON list of imports used in this chunk
+    return_type: Optional[str] = None  # Return type annotation
+
+    # Scoring
+    vector_score: float = 0.0  # Cosine similarity (0-1, higher is better)
+    bm25_score: float = 0.0    # BM25 score (unbounded, higher is better)
+    combined_score: float = 0.0  # Weighted combination
+
+    # Metadata
+    chunk_id: str = ""
+    milvus_id: Optional[int] = None
+    
+    def format_result(self, include_scores: bool = False, include_metadata: bool = False) -> str:
+        """Format the result for display"""
+        # Build header
+        header = f"📄 {self.file_name}"
+        if self.parent_context:
+            header += f" > {self.parent_context}"
+        if self.chunk_name:
+            header += f" > {self.chunk_name}"
+        if self.chunk_type:
+            header += f" ({self.chunk_type})"
+        header += f" [Lines {self.start_line}-{self.end_line}]"
+
+        # Add scores if requested
+        if include_scores:
+            header += f"\n   Scores: Vector={self.vector_score:.3f}, BM25={self.bm25_score:.3f}, Combined={self.combined_score:.3f}"
+
+        # Add enhanced metadata if available and requested
+        if include_metadata:
+            if self.signature:
+                header += f"\n   Signature: {self.signature}"
+            if self.return_type:
+                header += f"\n   Returns: {self.return_type}"
+            if self.docstring:
+                # Show first line of docstring
+                first_line = self.docstring.split('\n')[0][:80]
+                header += f"\n   Doc: {first_line}"
+
+        # Add content preview (first 3 lines)
+        lines = self.content.split('\n')[:3]
+        preview = '\n'.join(f"   {line}" for line in lines)
+        if len(self.content.split('\n')) > 3:
+            preview += "\n   ..."
+
+        return f"{header}\n{preview}"
+
+
+class HybridSearcher:
+    """
+    Hybrid search combining BM25 and vector search.
+    
+    This class provides:
+    1. BM25 for keyword/exact matching
+    2. Vector search for semantic similarity
+    3. Hybrid ranking combining both scores
+    """
+    
+    def __init__(
+        self,
+        config: ClaudeContextConfig,
+        embeddings: LocalEmbeddings,
+        milvus_manager: MilvusManager,
+        bm25_weight: float = 0.5,
+        vector_weight: float = 0.5
+    ):
+        """
+        Initialize the hybrid searcher.
+        
+        Args:
+            config: Configuration
+            embeddings: Embeddings model
+            milvus_manager: Milvus connection manager
+            bm25_weight: Weight for BM25 scores (0-1)
+            vector_weight: Weight for vector scores (0-1)
+        """
+        if not config:
+            raise ValueError("config is required")
+        if not embeddings:
+            raise ValueError("embeddings is required")
+        if not milvus_manager:
+            raise ValueError("milvus_manager is required")
+        
+        # Normalize weights
+        total_weight = bm25_weight + vector_weight
+        self.bm25_weight = bm25_weight / total_weight
+        self.vector_weight = vector_weight / total_weight
+        
+        self.config = config
+        self.embeddings = embeddings
+        self.milvus_client = milvus_manager.get_client()
+        self.collection_name = "code_chunks"
+        
+        # BM25 components (will be built from indexed data)
+        self.bm25_index: Optional[BM25Okapi] = None
+        self.bm25_documents: List[str] = []
+        self.bm25_metadata: List[Dict[str, Any]] = []
+        
+        logger.info(f"HybridSearcher initialized (BM25={self.bm25_weight:.2f}, Vector={self.vector_weight:.2f})")
+    
+    def build_bm25_index(self) -> None:
+        """
+        Build BM25 index from Milvus data.
+        
+        This loads all documents from Milvus and creates a BM25 index
+        for fast keyword searching.
+        """
+        logger.info("Building BM25 index from Milvus data...")
+        
+        # Check if collection exists
+        if not self.milvus_client.has_collection(self.collection_name):
+            raise ValueError(f"Collection '{self.collection_name}' does not exist. Index a repository first.")
+        
+        # Query all documents (without vectors to save memory)
+        # Note: In production, you might want to paginate this for large collections
+        results = self.milvus_client.query(
+            collection_name=self.collection_name,
+            filter="",  # No filter - get everything
+            output_fields=["content", "file_path", "file_name", "start_line", "end_line",
+                          "language", "chunk_type", "chunk_name", "parent_context", "chunk_id",
+                          "signature", "docstring", "imports", "return_type"],  # Enhanced fields
+            limit=10000  # Adjust based on your needs
+        )
+        
+        if not results:
+            logger.warning("No documents found in Milvus")
+            return
+        
+        # Prepare documents for BM25
+        self.bm25_documents = []
+        self.bm25_metadata = []
+        
+        for doc in results:
+            # Tokenize content for BM25
+            content = doc.get("content", "")
+            tokens = self._tokenize_code(content)
+            self.bm25_documents.append(tokens)
+            
+            # Store metadata for retrieval (including enhanced fields)
+            self.bm25_metadata.append({
+                "id": doc.get("id"),
+                "content": content,
+                "file_path": doc.get("file_path", ""),
+                "file_name": doc.get("file_name", ""),
+                "start_line": doc.get("start_line", 0),
+                "end_line": doc.get("end_line", 0),
+                "language": doc.get("language", ""),
+                "chunk_type": doc.get("chunk_type", ""),
+                "chunk_name": doc.get("chunk_name", ""),
+                "parent_context": doc.get("parent_context", ""),
+                "chunk_id": doc.get("chunk_id", ""),
+                # Enhanced fields
+                "signature": doc.get("signature", ""),
+                "docstring": doc.get("docstring", ""),
+                "imports": doc.get("imports", ""),
+                "return_type": doc.get("return_type", "")
+            })
+        
+        # Build BM25 index
+        self.bm25_index = BM25Okapi(self.bm25_documents)
+        
+        logger.info(f"✅ BM25 index built with {len(self.bm25_documents)} documents")
+    
+    def _tokenize_code(self, text: str) -> List[str]:
+        """
+        Tokenize code for BM25 indexing.
+        
+        This handles:
+        - CamelCase splitting (className -> class, Name)
+        - snake_case splitting (get_user -> get, user)
+        - Special characters removal
+        - Lowercasing for consistency
+        """
+        # Split on whitespace and special characters
+        tokens = re.findall(r'\b\w+\b', text.lower())
+        
+        # Further split camelCase and snake_case
+        expanded_tokens = []
+        for token in tokens:
+            # Split snake_case
+            parts = token.split('_')
+            for part in parts:
+                if part:
+                    # Split camelCase
+                    camel_parts = re.findall(r'[a-z]+|[A-Z][a-z]*', part)
+                    if camel_parts:
+                        expanded_tokens.extend(camel_parts)
+                    else:
+                        expanded_tokens.append(part)
+        
+        return expanded_tokens
+    
+    def search(
+        self,
+        query: str,
+        limit: int = 10,
+        filter_chunk_types: Optional[List[str]] = None,
+        filter_languages: Optional[List[str]] = None,
+        filter_paths: Optional[List[str]] = None,
+        rerank: bool = True
+    ) -> List[SearchResult]:
+        """
+        Perform hybrid search combining BM25 and vector search.
+
+        Args:
+            query: Search query
+            limit: Maximum number of results
+            filter_chunk_types: Filter by chunk types (function, class, etc.)
+            filter_languages: Filter by programming languages
+            filter_paths: Filter results to files under these directory prefixes
+            rerank: Whether to use hybrid reranking
+
+        Returns:
+            List of SearchResult objects sorted by relevance
+        """
+        if not query.strip():
+            logger.warning("Empty search query")
+            return []
+        
+        start_time = time.time()
+        
+        # Ensure BM25 index is built
+        if self.bm25_index is None:
+            self.build_bm25_index()
+            if self.bm25_index is None:
+                logger.warning("Could not build BM25 index, falling back to vector-only search")
+                rerank = False
+        
+        # Over-fetch when post-filters are active so we have enough after filtering
+        has_filters = filter_paths or filter_chunk_types or filter_languages
+        fetch_limit = limit * 5 if has_filters else limit * 2
+
+        # 1. Vector Search
+        vector_results = self._vector_search(query, fetch_limit)
+
+        # 2. BM25 Search (if available)
+        bm25_results = []
+        if self.bm25_index and rerank:
+            bm25_results = self._bm25_search(query, fetch_limit)
+
+        # 3. Combine and rerank
+        if rerank and bm25_results:
+            results = self._hybrid_rerank(vector_results, bm25_results, fetch_limit)
+        else:
+            results = vector_results[:fetch_limit]
+        
+        # 4. Apply filters
+        if filter_paths:
+            # Paths should be absolute; resolve any relative ones against cwd as fallback
+            resolved = [str(Path(p).resolve()) if not Path(p).is_absolute() else p
+                        for p in filter_paths]
+            results = [r for r in results
+                       if any(r.file_path.startswith(p) for p in resolved)]
+        if filter_chunk_types:
+            results = [r for r in results if r.chunk_type in filter_chunk_types]
+        if filter_languages:
+            results = [r for r in results if r.language in filter_languages]
+        
+        elapsed = time.time() - start_time
+        logger.info(f"Search completed in {elapsed:.3f}s - Found {len(results)} results")
+        
+        return results[:limit]
+    
+    def _vector_search(self, query: str, limit: int) -> List[SearchResult]:
+        """Perform vector similarity search"""
+        # Generate query embedding
+        query_embedding = self.embeddings.embed_texts([query])[0]
+
+        # Search in Milvus (including enhanced fields)
+        search_results = self.milvus_client.search(
+            collection_name=self.collection_name,
+            data=[query_embedding.tolist()],
+            limit=limit,
+            output_fields=["content", "file_path", "file_name", "start_line", "end_line",
+                          "language", "chunk_type", "chunk_name", "parent_context", "chunk_id",
+                          "signature", "docstring", "imports", "return_type"]
+        )
+
+        results = []
+        for hits in search_results:
+            for hit in hits:
+                # Create SearchResult with enhanced metadata
+                result = SearchResult(
+                    content=hit["entity"].get("content", ""),
+                    file_path=hit["entity"].get("file_path", ""),
+                    file_name=hit["entity"].get("file_name", ""),
+                    start_line=hit["entity"].get("start_line", 0),
+                    end_line=hit["entity"].get("end_line", 0),
+                    language=hit["entity"].get("language", ""),
+                    chunk_type=hit["entity"].get("chunk_type", ""),
+                    chunk_name=hit["entity"].get("chunk_name"),
+                    parent_context=hit["entity"].get("parent_context"),
+                    # Enhanced fields
+                    signature=hit["entity"].get("signature"),
+                    docstring=hit["entity"].get("docstring"),
+                    imports=hit["entity"].get("imports"),
+                    return_type=hit["entity"].get("return_type"),
+                    # Scoring
+                    vector_score=1 - hit["distance"],  # Convert distance to similarity
+                    chunk_id=hit["entity"].get("chunk_id", ""),
+                    milvus_id=hit.get("id")
+                )
+                results.append(result)
+
+        return results
+    
+    def _bm25_search(self, query: str, limit: int) -> List[SearchResult]:
+        """Perform BM25 keyword search"""
+        if not self.bm25_index:
+            return []
+        
+        # Tokenize query
+        query_tokens = self._tokenize_code(query)
+        
+        # Get BM25 scores
+        scores = self.bm25_index.get_scores(query_tokens)
+        
+        # Get top-k indices
+        top_indices = np.argsort(scores)[::-1][:limit]
+        
+        results = []
+        for idx in top_indices:
+            if scores[idx] > 0:  # Only include results with positive scores
+                metadata = self.bm25_metadata[idx]
+                result = SearchResult(
+                    content=metadata["content"],
+                    file_path=metadata["file_path"],
+                    file_name=metadata["file_name"],
+                    start_line=metadata["start_line"],
+                    end_line=metadata["end_line"],
+                    language=metadata["language"],
+                    chunk_type=metadata["chunk_type"],
+                    chunk_name=metadata["chunk_name"],
+                    parent_context=metadata["parent_context"],
+                    # Enhanced fields
+                    signature=metadata.get("signature"),
+                    docstring=metadata.get("docstring"),
+                    imports=metadata.get("imports"),
+                    return_type=metadata.get("return_type"),
+                    # Scoring
+                    bm25_score=float(scores[idx]),
+                    chunk_id=metadata["chunk_id"]
+                )
+                results.append(result)
+
+        return results
+    
+    def _hybrid_rerank(
+        self,
+        vector_results: List[SearchResult],
+        bm25_results: List[SearchResult],
+        limit: int
+    ) -> List[SearchResult]:
+        """
+        Combine and rerank results from vector and BM25 search.
+        
+        Uses reciprocal rank fusion with weighted scores.
+        """
+        # Create a map of all unique results
+        all_results: Dict[str, SearchResult] = {}
+        
+        # Add vector results
+        for i, result in enumerate(vector_results):
+            key = f"{result.file_path}:{result.start_line}"
+            if key not in all_results:
+                all_results[key] = result
+            else:
+                # Update vector score if we've seen this result
+                all_results[key].vector_score = max(
+                    all_results[key].vector_score,
+                    result.vector_score
+                )
+        
+        # Add/update BM25 results
+        for i, result in enumerate(bm25_results):
+            key = f"{result.file_path}:{result.start_line}"
+            if key in all_results:
+                # Update BM25 score for existing result
+                all_results[key].bm25_score = result.bm25_score
+            else:
+                # Add new result from BM25
+                all_results[key] = result
+        
+        # Normalize scores and calculate combined score
+        # Get max scores for normalization
+        max_vector = max((r.vector_score for r in all_results.values()), default=1.0)
+        max_bm25 = max((r.bm25_score for r in all_results.values()), default=1.0)
+        
+        for result in all_results.values():
+            # Normalize scores to 0-1 range
+            norm_vector = result.vector_score / max_vector if max_vector > 0 else 0
+            norm_bm25 = result.bm25_score / max_bm25 if max_bm25 > 0 else 0
+            
+            # Calculate weighted combined score
+            result.combined_score = (
+                self.vector_weight * norm_vector +
+                self.bm25_weight * norm_bm25
+            )
+        
+        # Sort by combined score
+        sorted_results = sorted(
+            all_results.values(),
+            key=lambda x: x.combined_score,
+            reverse=True
+        )
+        
+        return sorted_results[:limit]
+    
+    def explain_search(self, query: str, limit: int = 5) -> str:
+        """
+        Perform search and explain the scoring.
+        
+        Useful for debugging and understanding why certain results rank higher.
+        """
+        results = self.search(query, limit=limit)
+        
+        explanation = f"🔍 Search Query: '{query}'\n"
+        explanation += f"⚖️  Weights: BM25={self.bm25_weight:.2f}, Vector={self.vector_weight:.2f}\n\n"
+        
+        for i, result in enumerate(results, 1):
+            explanation += f"{i}. {result.format_result(include_scores=True)}\n\n"
+        
+        return explanation
+
+
+# Convenience function
+def create_hybrid_searcher(
+    config: ClaudeContextConfig,
+    embeddings: LocalEmbeddings,
+    milvus_manager: MilvusManager,
+    bm25_weight: float = 0.5
+) -> HybridSearcher:
+    """
+    Create a hybrid searcher with balanced weights.
+    
+    Args:
+        config: Configuration
+        embeddings: Embeddings model
+        milvus_manager: Milvus manager
+        bm25_weight: Weight for BM25 (0-1), vector weight will be 1-bm25_weight
+        
+    Returns:
+        Configured HybridSearcher
+    """
+    vector_weight = 1.0 - bm25_weight
+    return HybridSearcher(
+        config=config,
+        embeddings=embeddings,
+        milvus_manager=milvus_manager,
+        bm25_weight=bm25_weight,
+        vector_weight=vector_weight
+    )

claude_context/skills/__init__.py

@@ -0,0 +1,11 @@
+"""
+Bridge skills for vibe2doc.
+
+Standalone CLI-callable modules that expose vibe2doc's code analysis
+capabilities for use in any documentation workflow.
+
+Skills:
+    evidence_retrieval: Search codebase for evidence matching a query
+    grounded_review: Verify document claims against source code
+    api_surface: Extract API surface from source files (deterministic)
+"""

claude_context/skills/_cli_common.py

@@ -0,0 +1,74 @@
+"""
+Shared CLI utilities for bridge skills.
+
+Handles argument parsing, logging to stderr, and JSON output routing
+so that stdout stays clean for machine-consumable JSON.
+"""
+
+import argparse
+import json
+import logging
+import sys
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+def setup_logging():
+    """Configure logging to stderr so stdout stays clean for JSON output."""
+    root = logging.getLogger()
+    root.handlers.clear()
+    handler = logging.StreamHandler(sys.stderr)
+    handler.setFormatter(logging.Formatter(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    ))
+    root.addHandler(handler)
+    root.setLevel(logging.INFO)
+
+
+def add_common_args(parser: argparse.ArgumentParser):
+    """Add shared arguments for dual-mode CLI (standalone + docs-orchestrator)."""
+    parser.add_argument(
+        '--ticket', type=str, default=None,
+        help='Docs-orchestrator ticket ID (stored in output metadata)'
+    )
+    parser.add_argument(
+        '--base-path', type=str, default=None,
+        help='Docs-orchestrator output base path. When set, writes JSON to file instead of stdout.'
+    )
+    parser.add_argument(
+        '--db-path', type=str, default=None,
+        help='Override Milvus DB path (default: {repo}/.vibe2doc/index.db)'
+    )
+    parser.add_argument(
+        '--reindex', action='store_true',
+        help='Force re-indexing even if an existing index is found'
+    )
+
+
+def output_result(result: Dict[str, Any], args: argparse.Namespace, filename: str = "output.json"):
+    """Output JSON result to stdout or file depending on mode.
+
+    Args:
+        result: The result dict to serialize
+        args: Parsed CLI args (checks for --base-path and --ticket)
+        filename: Output filename when writing to --base-path
+    """
+    if args.ticket:
+        result["ticket"] = args.ticket
+
+    json_str = json.dumps(result, indent=2, default=str)
+
+    if args.base_path:
+        out_dir = Path(args.base_path)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        out_path = out_dir / filename
+        out_path.write_text(json_str)
+        logging.getLogger(__name__).info(f"Output written to {out_path}")
+    else:
+        print(json_str)
+
+
+def error_result(error: str, exit_code: int = 1):
+    """Print a JSON error object to stdout and exit."""
+    print(json.dumps({"error": error, "exit_code": exit_code}))
+    sys.exit(exit_code)

claude_context/skills/_index_manager.py

@@ -0,0 +1,98 @@
+"""
+Index lifecycle manager for bridge skills.
+
+Handles the expensive indexing operation with reuse: if an existing
+Milvus DB is found for a repo, reuse it. Otherwise, index from scratch.
+"""
+
+import logging
+from pathlib import Path
+from typing import Dict, Any, Optional, Tuple
+
+from ..config import ClaudeContextConfig, MilvusManager
+from ..embeddings import LocalEmbeddings
+from ..indexer import RepositoryIndexer
+from ..search import HybridSearcher, create_hybrid_searcher
+
+logger = logging.getLogger(__name__)
+
+
+def _default_db_path(repo_path: str) -> str:
+    """Compute the default index DB path co-located with the repo."""
+    return str(Path(repo_path).resolve() / ".vibe2doc" / "index.db")
+
+
+def ensure_index(
+    repo_path: str,
+    db_path: Optional[str] = None,
+    reindex: bool = False,
+) -> Tuple[HybridSearcher, Dict[str, Any]]:
+    """Ensure a searchable index exists for the given repo.
+
+    If an index already exists and reindex=False, reuses it.
+    Otherwise, runs full indexing (discover, chunk, embed, store).
+
+    Args:
+        repo_path: Path to the repository to index
+        db_path: Override Milvus DB path (default: {repo}/.vibe2doc/index.db)
+        reindex: Force re-indexing even if index exists
+
+    Returns:
+        (searcher, info) where info contains index metadata
+    """
+    repo_path = str(Path(repo_path).resolve())
+    if db_path is None:
+        db_path = _default_db_path(repo_path)
+
+    # Ensure the DB directory exists
+    Path(db_path).parent.mkdir(parents=True, exist_ok=True)
+
+    db_exists = Path(db_path).exists()
+    config = ClaudeContextConfig(db_path=db_path)
+    milvus_mgr = MilvusManager(config)
+    embeddings = LocalEmbeddings(config.embedding_model)
+
+    info = {
+        "repo_path": repo_path,
+        "index_path": db_path,
+        "reused_existing": False,
+        "chunks_count": 0,
+    }
+
+    if db_exists and not reindex:
+        # Reuse existing index
+        logger.info(f"Reusing existing index at {db_path}")
+        client = milvus_mgr.get_client()
+        if client.has_collection(config.collection_name):
+            count = client.query(
+                collection_name=config.collection_name,
+                filter="",
+                output_fields=["count(*)"],
+            )
+            info["chunks_count"] = count[0].get("count(*)", 0) if count else 0
+            info["reused_existing"] = True
+        else:
+            # DB file exists but collection doesn't — need to index
+            logger.info("DB exists but collection missing, indexing from scratch")
+            indexer = RepositoryIndexer(config, embeddings, milvus_mgr)
+            stats = indexer.index_repository(repo_path, show_progress=True)
+            info["chunks_count"] = stats.get("chunks_indexed", 0)
+    else:
+        # Fresh index
+        if reindex:
+            logger.info(f"Re-indexing {repo_path} (forced)")
+        else:
+            logger.info(f"No existing index found, indexing {repo_path}")
+        indexer = RepositoryIndexer(config, embeddings, milvus_mgr)
+        stats = indexer.index_repository(repo_path, show_progress=True)
+        info["chunks_count"] = stats.get("chunks_indexed", 0)
+
+    # Create searcher and build BM25 index
+    searcher = create_hybrid_searcher(config, embeddings, milvus_mgr)
+    searcher.build_bm25_index()
+
+    logger.info(
+        f"Index ready: {info['chunks_count']} chunks, "
+        f"reused={info['reused_existing']}"
+    )
+    return searcher, info