contextual-engine 0.1.0__py3-none-any.whl

contextual/retrieval/search.py

@@ -0,0 +1,548 @@
+"""Hybrid search engine: BM25 + vector similarity with Reciprocal Rank Fusion.
+
+Implements the search pipeline per TECHNICAL_SPEC §5:
+- BM25 via FTS5 (top-100 candidates)
+- Vector cosine similarity via vec0 (top-100 candidates)
+- RRF fusion with configurable weights (BM25=0.6, Dense=0.4)
+- Returns ranked results with scores
+
+This module is the retrieval workhorse - all MCP search queries flow through here.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from contextual.core.errors import ErrorCode, RetrievalError, retrieval_context
+from contextual.observability.logging import get_logger
+from contextual.storage.vec0_manager import search_vectors
+
+
+if TYPE_CHECKING:
+    from sqlite3 import Connection
+
+
+logger = get_logger(__name__)
+
+# ============================================================================
+# CONSTANTS (from TECHNICAL_SPEC §5)
+# ============================================================================
+
+BM25_CANDIDATES = 100
+"""Number of candidates to fetch from BM25 index."""
+
+DENSE_CANDIDATES = 100
+"""Number of candidates to fetch from vector search."""
+
+RRF_K = 60
+"""Reciprocal Rank Fusion constant (Cormack et al. 2009)."""
+
+RRF_BM25_WEIGHT = 0.6
+"""Weight for BM25 results in RRF scoring (identifiers dominate code queries)."""
+
+RRF_DENSE_WEIGHT = 0.4
+"""Weight for dense vector results in RRF scoring."""
+
+
+# ============================================================================
+# DATA MODELS
+# ============================================================================
+
+
+@dataclass(frozen=True)
+class SearchResult:
+    """A single search result with metadata and score.
+
+    Attributes:
+        chunk_id: Database chunk ID.
+        path: File path relative to project root.
+        symbol_name: Function/class name (or None for module chunks).
+        chunk_type: Type of code structure.
+        start_line: First line of chunk.
+        end_line: Last line of chunk.
+        score: Combined relevance score (0.0-1.0, higher = more relevant).
+        bm25_rank: Rank in BM25 results (None if not in BM25 top-100).
+        dense_rank: Rank in vector results (None if not in vector top-100).
+        snippet: Preview of chunk content (first 200 chars).
+    """
+
+    chunk_id: int
+    path: str
+    symbol_name: str | None
+    chunk_type: str
+    start_line: int
+    end_line: int
+    score: float
+    bm25_rank: int | None
+    dense_rank: int | None
+    snippet: str
+
+
+# ============================================================================
+# BM25 SEARCH (FTS5)
+# ============================================================================
+
+
+def _bm25_search(
+    conn: Connection,
+    query: str,
+    k: int = BM25_CANDIDATES,
+    language: str | None = None,
+) -> list[tuple[int, float, int]]:
+    """Execute BM25 full-text search via FTS5.
+
+    Args:
+        conn: SQLite connection.
+        query: Search query string (supports FTS5 syntax).
+        k: Number of results to return.
+        language: Optional language filter (e.g., 'python', 'typescript').
+
+    Returns:
+        List of (chunk_id, bm25_score, rank) tuples ordered by descending score.
+
+    Raises:
+        RetrievalError: If FTS5 query fails.
+    """
+    try:
+        cursor = conn.cursor()
+
+        # Extract keywords from natural language query
+        # Remove common stop words and keep only meaningful terms
+        stop_words = {
+            'a', 'an', 'the', 'is', 'are', 'was', 'were', 'be', 'been',
+            'how', 'what', 'when', 'where', 'why', 'which', 'who',
+            'does', 'do', 'did', 'can', 'could', 'should', 'would',
+            'to', 'of', 'in', 'on', 'at', 'for', 'with', 'from', 'by',
+            'i', 'you', 'we', 'they', 'it', 'this', 'that', 'these', 'those',
+        }
+        
+        # Tokenize and filter
+        tokens = query.lower().split()
+        keywords = [t for t in tokens if t not in stop_words and len(t) > 2]
+        
+        # Join keywords with OR operator for FTS5
+        fts_query = ' OR '.join(keywords) if keywords else query.strip()
+        
+        # Preserve exact symbols with underscores (snake_case identifiers)
+        # For single-word queries containing underscores, force phrase match
+        if '_' in query and len(query.split()) == 1:
+            fts_query = f'"{query}"'  # Force phrase match for snake_case identifiers
+        
+        logger.debug("BM25 query processing", original=query, keywords=keywords, fts_query=fts_query)
+
+        # Build query parameters
+        params: list[str | int] = [fts_query]
+
+        # Add language filter if specified
+        language_filter = ""
+        if language:
+            language_filter = "AND f.language = ?"
+            params.append(language)
+
+        # Query FTS5 with BM25 ranking (C4: Boost symbol_name with 1.0, 2.5, 0.0)
+        # Note: Use subquery to compute bm25() scores before window function
+        # (FTS5 bm25() cannot be used directly inside window function context)
+        sql = f"""
+            SELECT 
+                chunk_id,
+                bm25_score,
+                ROW_NUMBER() OVER (ORDER BY bm25_score DESC) AS rank
+            FROM (
+                SELECT 
+                    c.id as chunk_id,
+                    -bm25(fts_chunks) AS bm25_score
+                FROM fts_chunks fts
+                JOIN chunks c ON fts.rowid = c.id
+                JOIN files f ON c.file_id = f.id
+                WHERE fts_chunks MATCH ?
+                {language_filter}
+            )
+            ORDER BY bm25_score DESC
+            LIMIT ?
+        """
+        params.append(k)
+
+        cursor.execute(sql, tuple(params))
+        results = cursor.fetchall()
+
+        logger.debug(
+            "BM25 search completed",
+            query=query,
+            results_count=len(results),
+            language=language,
+        )
+
+        return results
+
+    except sqlite3.Error as e:
+        msg = f"BM25 search failed: {e}"
+        raise RetrievalError(
+            code=ErrorCode.RETRIEVAL_BM25_FAILED,
+            message=msg,
+            context=retrieval_context(
+                query=query,
+            ),
+        ) from e
+
+
+# ============================================================================
+# VECTOR SEARCH (vec0)
+# ============================================================================
+
+
+def _vector_search(
+    conn: Connection,
+    query_embedding: list[float],
+    k: int = DENSE_CANDIDATES,
+    language: str | None = None,
+) -> list[tuple[int, float, int]]:
+    """Execute vector similarity search via vec0.
+
+    Args:
+        conn: SQLite connection.
+        query_embedding: Query vector (768-dim).
+        k: Number of results to return.
+        language: Optional language filter.
+
+    Returns:
+        List of (chunk_id, similarity_score, rank) tuples ordered by descending score.
+
+    Raises:
+        RetrievalError: If vector search fails.
+    """
+    try:
+        # Get raw vector results (chunk_id, score)
+        raw_results = search_vectors(conn, query_embedding, k=k)
+
+        # If language filter specified, fetch chunk metadata and filter
+        if language:
+            cursor = conn.cursor()
+            chunk_ids = [chunk_id for chunk_id, _ in raw_results]
+
+            if not chunk_ids:
+                return []
+
+            # Fetch chunks with language filter
+            placeholders = ",".join("?" * len(chunk_ids))
+            sql = f"""
+                SELECT id
+                FROM chunks
+                WHERE id IN ({placeholders})
+                AND language = ?
+            """
+
+            cursor.execute(sql, (*chunk_ids, language))
+            valid_ids = {row[0] for row in cursor.fetchall()}
+
+            # Filter results to only valid chunk IDs
+            filtered_results = [
+                (chunk_id, score)
+                for chunk_id, score in raw_results
+                if chunk_id in valid_ids
+            ]
+
+            results = filtered_results[:k]
+        else:
+            results = raw_results
+
+        # Add rank (1-indexed)
+        ranked_results = [
+            (chunk_id, score, rank + 1)
+            for rank, (chunk_id, score) in enumerate(results)
+        ]
+
+        logger.debug(
+            "Vector search completed",
+            results_count=len(ranked_results),
+            language=language,
+        )
+
+        return ranked_results
+
+    except Exception as e:
+        msg = f"Vector search failed: {e}"
+        raise RetrievalError(
+            code=ErrorCode.RETRIEVAL_DENSE_SEARCH_FAILED,
+            message=msg,
+            context=retrieval_context(),
+        ) from e
+
+
+# ============================================================================
+# RECIPROCAL RANK FUSION
+# ============================================================================
+
+
+def _reciprocal_rank_fusion(
+    bm25_results: list[tuple[int, float, int]],
+    vector_results: list[tuple[int, float, int]],
+    k: int = RRF_K,
+    bm25_weight: float = RRF_BM25_WEIGHT,
+    dense_weight: float = RRF_DENSE_WEIGHT,
+) -> list[tuple[int, float, int | None, int | None]]:
+    """Fuse BM25 and vector results using Reciprocal Rank Fusion.
+
+    RRF formula: score(d) = Σ_s [ weight_s / (k + rank_s(d)) ]
+    where s ∈ {bm25, dense}, weight_s is the source weight, rank_s(d)
+    is the rank of document d in source s.
+
+    Args:
+        bm25_results: BM25 results as (chunk_id, score, rank).
+        vector_results: Vector results as (chunk_id, score, rank).
+        k: RRF constant (higher = less emphasis on top ranks).
+        bm25_weight: Weight for BM25 results.
+        dense_weight: Weight for vector results.
+
+    Returns:
+        List of (chunk_id, rrf_score, bm25_rank, dense_rank) tuples
+        ordered by descending RRF score.
+    """
+    # Build rank maps: chunk_id -> rank
+    bm25_ranks = {chunk_id: rank for chunk_id, _, rank in bm25_results}
+    dense_ranks = {chunk_id: rank for chunk_id, _, rank in vector_results}
+
+    # Collect all unique chunk IDs from both sources
+    all_chunk_ids = set(bm25_ranks.keys()) | set(dense_ranks.keys())
+
+    # Compute RRF scores
+    rrf_scores: dict[int, float] = {}
+    for chunk_id in all_chunk_ids:
+        score = 0.0
+
+        if chunk_id in bm25_ranks:
+            score += bm25_weight / (k + bm25_ranks[chunk_id])
+
+        if chunk_id in dense_ranks:
+            score += dense_weight / (k + dense_ranks[chunk_id])
+
+        rrf_scores[chunk_id] = score
+
+    # Sort by descending RRF score
+    sorted_results = sorted(
+        rrf_scores.items(),
+        key=lambda x: x[1],
+        reverse=True,
+    )
+
+    # Build final result list with source ranks
+    fused_results = [
+        (
+            chunk_id,
+            score,
+            bm25_ranks.get(chunk_id),
+            dense_ranks.get(chunk_id),
+        )
+        for chunk_id, score in sorted_results
+    ]
+
+    logger.debug(
+        "RRF fusion completed",
+        total_candidates=len(all_chunk_ids),
+        bm25_only=len(bm25_ranks - dense_ranks.keys()),
+        dense_only=len(dense_ranks - bm25_ranks.keys()),
+        overlap=len(bm25_ranks.keys() & dense_ranks.keys()),
+    )
+
+    return fused_results
+
+
+# ============================================================================
+# RESULT HYDRATION
+# ============================================================================
+
+
+def _hydrate_results(
+    conn: Connection,
+    fused_results: list[tuple[int, float, int | None, int | None]],
+    top_k: int,
+) -> list[SearchResult]:
+    """Fetch chunk metadata and build SearchResult objects.
+
+    Args:
+        conn: SQLite connection.
+        fused_results: RRF results as (chunk_id, score, bm25_rank, dense_rank).
+        top_k: Number of results to return.
+
+    Returns:
+        List of SearchResult objects with full metadata.
+
+    Raises:
+        RetrievalError: If metadata fetch fails.
+    """
+    if not fused_results:
+        return []
+
+    # Take top-k
+    top_results = fused_results[:top_k]
+    chunk_ids = [chunk_id for chunk_id, _, _, _ in top_results]
+
+    try:
+        cursor = conn.cursor()
+
+        # Fetch chunk metadata
+        placeholders = ",".join("?" * len(chunk_ids))
+        sql = f"""
+            SELECT 
+                c.id,
+                f.path,
+                c.symbol_name,
+                c.symbol_type,
+                c.start_line,
+                c.end_line,
+                c.content
+            FROM chunks c
+            JOIN files f ON c.file_id = f.id
+            WHERE c.id IN ({placeholders})
+        """
+
+        cursor.execute(sql, chunk_ids)
+        rows = cursor.fetchall()
+
+        # Build chunk metadata map
+        chunk_meta = {
+            row[0]: {
+                "path": row[1],
+                "symbol_name": row[2],
+                "chunk_type": row[3],
+                "start_line": row[4],
+                "end_line": row[5],
+                "content": row[6],
+            }
+            for row in rows
+        }
+
+        # Build SearchResult objects preserving RRF order
+        results = []
+        for chunk_id, score, bm25_rank, dense_rank in top_results:
+            if chunk_id not in chunk_meta:
+                # Chunk was deleted between search and hydration - skip it
+                logger.warning("Chunk not found during hydration", chunk_id=chunk_id)
+                continue
+
+            meta = chunk_meta[chunk_id]
+
+            # Generate snippet (first 200 chars)
+            content = meta["content"]
+            snippet = content[:200] + "..." if len(content) > 200 else content
+
+            result = SearchResult(
+                chunk_id=chunk_id,
+                path=meta["path"],
+                symbol_name=meta["symbol_name"],
+                chunk_type=meta["chunk_type"],
+                start_line=meta["start_line"],
+                end_line=meta["end_line"],
+                score=score,
+                bm25_rank=bm25_rank,
+                dense_rank=dense_rank,
+                snippet=snippet,
+            )
+
+            results.append(result)
+
+        logger.debug(
+            "Results hydrated",
+            requested=top_k,
+            returned=len(results),
+        )
+
+        return results
+
+    except sqlite3.Error as e:
+        msg = f"Failed to hydrate search results: {e}"
+        raise RetrievalError(
+            code=ErrorCode.RETRIEVAL_RERANK_FAILED,  # Using existing code for now
+            message=msg,
+            context=retrieval_context(),
+        ) from e
+
+
+# ============================================================================
+# PUBLIC API
+# ============================================================================
+
+
+def hybrid_search(
+    conn: Connection,
+    query: str,
+    query_embedding: list[float],
+    top_k: int = 10,
+    language: str | None = None,
+    bm25_weight: float | None = None,
+    dense_weight: float | None = None,
+) -> list[SearchResult]:
+    """Execute hybrid BM25 + vector search with RRF fusion.
+
+    This is the main search entry point. It:
+    1. Runs BM25 search (top-100)
+    2. Runs vector search (top-100)
+    3. Fuses results with RRF
+    4. Hydrates metadata
+    5. Returns top-k results
+
+    Args:
+        conn: SQLite connection (must have vec0 loaded).
+        query: Natural language query string.
+        query_embedding: Query embedding vector (768-dim).
+        top_k: Number of results to return (1-100).
+        language: Optional language filter.
+        bm25_weight: Optional BM25 weight for RRF (default: RRF_BM25_WEIGHT).
+                    Source-type aware: pass 0.55 for code, 0.4 for docs.
+        dense_weight: Optional dense weight for RRF (default: RRF_DENSE_WEIGHT).
+                     Source-type aware: pass 0.45 for code, 0.6 for docs.
+
+    Returns:
+        List of SearchResult objects ordered by relevance.
+
+    Raises:
+        RetrievalError: If search fails at any stage.
+
+    Example:
+        >>> from contextual.embedding import get_embedder
+        >>> embedder = get_embedder()
+        >>> query = "parse JSON configuration"
+        >>> embedding = embedder.embed([query])[0]
+        >>> results = hybrid_search(conn, query, embedding, top_k=10)
+        >>> for r in results:
+        ...     print(f"{r.path}:{r.start_line} - {r.score:.3f}")
+    """
+    # Use provided weights or fall back to defaults
+    _bm25_weight = bm25_weight if bm25_weight is not None else RRF_BM25_WEIGHT
+    _dense_weight = dense_weight if dense_weight is not None else RRF_DENSE_WEIGHT
+
+    logger.info(
+        "Starting hybrid search",
+        query=query,
+        top_k=top_k,
+        language=language,
+        bm25_weight=_bm25_weight,
+        dense_weight=_dense_weight,
+    )
+
+    # Execute both searches in parallel (SQLite WAL allows concurrent reads)
+    bm25_results = _bm25_search(conn, query, k=BM25_CANDIDATES, language=language)
+    vector_results = _vector_search(conn, query_embedding, k=DENSE_CANDIDATES, language=language)
+
+    # Fuse with RRF using source-type-aware weights
+    fused_results = _reciprocal_rank_fusion(
+        bm25_results,
+        vector_results,
+        k=RRF_K,
+        bm25_weight=_bm25_weight,
+        dense_weight=_dense_weight,
+    )
+
+    # Hydrate and return
+    results = _hydrate_results(conn, fused_results, top_k)
+
+    logger.info(
+        "Hybrid search completed",
+        query=query,
+        results_count=len(results),
+        bm25_candidates=len(bm25_results),
+        vector_candidates=len(vector_results),
+    )
+
+    return results

contextual/security/__init__.py

@@ -0,0 +1,52 @@
+"""Security hardening: path safety, input sanitization, and workspace isolation.
+
+Prevents path traversal, SQL/FTS5 injection, Unicode-based attacks, and prompt
+injection via indexed content. Implements per-project workspace isolation.
+"""
+from __future__ import annotations
+
+from contextual.security.paths import (
+    ensure_parent_exists,
+    is_safe_path,
+    safe_join,
+    safe_path,
+    validate_filename,
+)
+from contextual.security.sanitize import (
+    detect_prompt_injection_patterns,
+    escape_for_logging,
+    escape_fts5_match_query,
+    sanitize_file_content,
+    sanitize_user_input,
+    strip_unicode_control_chars,
+    validate_sql_identifier,
+    wrap_untrusted_content,
+)
+from contextual.security.workspace import (
+    WorkspaceManager,
+    get_workspace_manager,
+    init_workspace,
+)
+
+
+__all__ = [
+    # Path safety
+    "safe_path",
+    "is_safe_path",
+    "safe_join",
+    "validate_filename",
+    "ensure_parent_exists",
+    # Input sanitization
+    "strip_unicode_control_chars",
+    "escape_fts5_match_query",
+    "validate_sql_identifier",
+    "sanitize_user_input",
+    "sanitize_file_content",
+    "detect_prompt_injection_patterns",
+    "wrap_untrusted_content",
+    "escape_for_logging",
+    # Workspace management
+    "WorkspaceManager",
+    "init_workspace",
+    "get_workspace_manager",
+]