code-review-graph-codeblackwell 2.3.6.post1__py3-none-any.whl

code_review_graph/search.py

@@ -0,0 +1,447 @@
+"""Hybrid search engine combining FTS5 (BM25) and vector embeddings.
+
+Uses Reciprocal Rank Fusion (RRF) to merge results from full-text search
+and semantic similarity, with query-aware kind boosting and context-file
+boosting for relevance tuning.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import sqlite3
+from typing import Any, Optional
+
+from .graph import GraphStore, _sanitize_name
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# FTS5 index management
+# ---------------------------------------------------------------------------
+
+
+def rebuild_fts_index(store: GraphStore) -> int:
+    """Rebuild the FTS5 index from the nodes table.
+
+    Checks whether the ``nodes_fts`` virtual table exists, clears it, then
+    repopulates it from every row in ``nodes``.
+
+    Returns:
+        Number of rows indexed.
+    """
+    # NOTE: rebuild_fts_index uses store._conn directly because it manages
+    # the FTS5 virtual table DDL, which is tightly coupled to SQLite internals.
+    conn = store._conn
+
+    # Wrap the full DROP + CREATE + INSERT sequence in an explicit transaction
+    # so a crash mid-rebuild cannot leave the DB without an FTS table at all
+    # (DROP succeeded but CREATE/INSERT didn't).  See #259.
+    if conn.in_transaction:
+        logger.warning("Rolling back uncommitted transaction before BEGIN IMMEDIATE")
+        conn.rollback()
+    conn.execute("BEGIN IMMEDIATE")
+    try:
+        # Drop and recreate the FTS table with content sync to match migration v5
+        conn.execute("DROP TABLE IF EXISTS nodes_fts")
+        conn.execute("""
+            CREATE VIRTUAL TABLE nodes_fts USING fts5(
+                name, qualified_name, file_path, signature,
+                content='nodes', content_rowid='rowid',
+                tokenize='porter unicode61'
+            )
+        """)
+
+        # Rebuild from the content table (nodes) using the FTS5 rebuild command
+        conn.execute("INSERT INTO nodes_fts(nodes_fts) VALUES('rebuild')")
+
+
+        conn.commit()
+    except BaseException:
+        conn.rollback()
+        raise
+
+    count = conn.execute("SELECT count(*) FROM nodes_fts").fetchone()[0]
+    logger.info("FTS index rebuilt: %d rows indexed", count)
+    return count
+
+
+# ---------------------------------------------------------------------------
+# Query kind boosting heuristics
+# ---------------------------------------------------------------------------
+
+
+_DOTTED_IDENT_RE = re.compile(r'\b[A-Za-z_][\w]*(?:\.[A-Za-z_][\w]*)+\b')
+_SNAKE_IDENT_RE = re.compile(r'\b[a-z][a-z0-9]*(?:_[a-z0-9]+)+\b')
+_PASCAL_IDENT_RE = re.compile(r'\b[A-Z][a-z0-9]+(?:[A-Z][a-z0-9]+)+\b')
+
+
+def extract_query_identifiers(query: str) -> list[str]:
+    """Pull out identifier-shaped tokens from anywhere in a query.
+
+    Catches dotted forms (``Context.Next``), snake_case (``get_dependant``),
+    and CamelCase (``APIRoute``) even when they're embedded in a natural-
+    language sentence. Used to boost search hits whose qualified_name
+    contains any of these tokens, so an LLM asking "Who advances the gin
+    middleware chain via Context.Next" lands on ``Context.Next`` instead of
+    the bare ``Context`` class.
+    """
+    found: list[str] = []
+    seen: set[str] = set()
+    for pat in (_DOTTED_IDENT_RE, _SNAKE_IDENT_RE, _PASCAL_IDENT_RE):
+        for match in pat.findall(query):
+            lo = match.lower()
+            if lo not in seen and len(lo) >= 3:
+                seen.add(lo)
+                found.append(lo)
+    return found
+
+
+def detect_query_kind_boost(query: str) -> dict[str, Any]:
+    """Detect query patterns and return per-node boost multipliers.
+
+    Heuristics:
+    - PascalCase queries (e.g. ``MyClass``) boost Class/Type by 1.5x
+    - snake_case queries (e.g. ``get_users``) boost Function by 1.5x
+    - Queries containing ``.`` boost qualified name matches by 2.0x
+    - Identifier-shaped tokens *anywhere* in the query (dotted, snake_case,
+      CamelCase) boost results whose qualified_name contains them by 2.0x.
+      See ``extract_query_identifiers``.
+
+    Returns:
+        Dict whose keys are either node kind strings (mapped to float
+        multipliers) or one of the special keys ``_qualified``,
+        ``_qualified_identifiers``.
+    """
+    boosts: dict[str, Any] = {}
+
+    if not query or not query.strip():
+        return boosts
+
+    q = query.strip()
+
+    # PascalCase: starts with uppercase, has at least one lowercase after
+    if re.match(r'^[A-Z][a-z]', q) and not q.isupper():
+        boosts["Class"] = 1.5
+        boosts["Type"] = 1.5
+
+    # snake_case or SCREAMING_SNAKE_CASE: contains underscore with letters
+    if '_' in q and re.search(r'[a-zA-Z]', q):
+        boosts["Function"] = 1.5
+
+    # Dotted path: boost qualified name matches
+    if '.' in q:
+        boosts["_qualified"] = 2.0
+
+    # Identifiers extracted from anywhere in the query
+    idents = extract_query_identifiers(q)
+    if idents:
+        boosts["_qualified_identifiers"] = idents
+
+    return boosts
+
+
+# ---------------------------------------------------------------------------
+# Reciprocal Rank Fusion
+# ---------------------------------------------------------------------------
+
+
+def rrf_merge(*result_lists: list[tuple[int, float]], k: int = 60) -> list[tuple[int, float]]:
+    """Merge multiple ranked result lists using Reciprocal Rank Fusion.
+
+    Each input list contains ``(id, score)`` tuples, ordered by score
+    descending. The RRF score for each item is the sum of
+    ``1 / (k + rank + 1)`` across all lists it appears in, where rank is
+    the 0-based position.
+
+    Args:
+        *result_lists: Variable number of ranked result lists.
+        k: RRF constant (default 60). Higher values reduce the impact of
+           rank differences.
+
+    Returns:
+        Merged list of ``(id, rrf_score)`` tuples sorted by score descending.
+    """
+    scores: dict[int, float] = {}
+
+    for result_list in result_lists:
+        for rank, (item_id, _score) in enumerate(result_list):
+            scores[item_id] = scores.get(item_id, 0.0) + 1.0 / (k + rank + 1)
+
+    merged = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+    return merged
+
+
+# ---------------------------------------------------------------------------
+# FTS5 search
+# ---------------------------------------------------------------------------
+
+
+def _fts_search(
+    conn: sqlite3.Connection,
+    query: str,
+    limit: int = 50,
+) -> list[tuple[int, float]]:
+    """Run an FTS5 BM25 search against the nodes_fts table.
+
+    Returns list of ``(node_id, bm25_score)`` tuples. The BM25 score is
+    negated so higher = better (FTS5 returns negative BM25).
+    """
+    # Sanitize: wrap in double quotes to prevent FTS5 operator injection
+    safe_query = '"' + query.replace('"', '""') + '"'
+
+    try:
+        rows = conn.execute(
+            "SELECT rowid, rank FROM nodes_fts WHERE nodes_fts MATCH ? "
+            "ORDER BY rank LIMIT ?",
+            (safe_query, limit),
+        ).fetchall()
+        # FTS5 rank is negative BM25 (lower = better), negate for consistency
+        return [(row[0], -row[1]) for row in rows]
+    except sqlite3.OperationalError as e:
+        logger.warning("FTS5 search failed: %s", e)
+        return []
+
+
+# ---------------------------------------------------------------------------
+# Embedding search (optional)
+# ---------------------------------------------------------------------------
+
+
+def _embedding_search(
+    store: GraphStore,
+    query: str,
+    limit: int = 50,
+    model: str | None = None,
+    provider: str | None = None,
+) -> list[tuple[int, float]]:
+    """Run a vector similarity search using the embedding store.
+
+    Returns list of ``(node_id, similarity_score)`` tuples.
+    Gracefully returns an empty list if embeddings are not available.
+    """
+    try:
+        from .embeddings import EmbeddingStore
+    except ImportError:
+        return []
+
+    try:
+        emb_store = EmbeddingStore(store.db_path, provider=provider, model=model)
+        try:
+            if not emb_store.available or emb_store.count() == 0:
+                return []
+
+            results = emb_store.search(query, limit=limit)
+            # Map qualified names back to node IDs
+            id_scores: list[tuple[int, float]] = []
+            for qn, score in results:
+                node = store.get_node(qn)
+                if node:
+                    id_scores.append((node.id, score))
+            return id_scores
+        finally:
+            emb_store.close()
+    except Exception as e:
+        logger.warning("Embedding search failed: %s", e)
+        return []
+
+
+# ---------------------------------------------------------------------------
+# Keyword LIKE fallback
+# ---------------------------------------------------------------------------
+
+
+def _keyword_search(
+    conn: sqlite3.Connection,
+    query: str,
+    limit: int = 50,
+) -> list[tuple[int, float]]:
+    """Fall back to simple LIKE keyword matching.
+
+    Each word in the query must match independently (AND logic).
+    Returns ``(node_id, score)`` tuples with a basic relevance score.
+    """
+    words = query.lower().split()
+    if not words:
+        return []
+
+    conditions: list[str] = []
+    params: list[str | int] = []
+    for word in words:
+        conditions.append(
+            "(LOWER(name) LIKE ? OR LOWER(qualified_name) LIKE ?)"
+        )
+        params.extend([f"%{word}%", f"%{word}%"])
+
+    where = " AND ".join(conditions)
+    params.append(limit)
+    sql = f"SELECT id, name, qualified_name FROM nodes WHERE {where} LIMIT ?"  # nosec B608
+
+    try:
+        rows = conn.execute(sql, params).fetchall()
+    except sqlite3.OperationalError:
+        return []
+
+    # Assign a simple relevance score: exact name match > prefix > contains
+    q_lower = query.lower()
+    results: list[tuple[int, float]] = []
+    for row in rows:
+        name_lower = row["name"].lower()
+        if name_lower == q_lower:
+            score = 3.0
+        elif name_lower.startswith(q_lower):
+            score = 2.0
+        else:
+            score = 1.0
+        results.append((row["id"], score))
+
+    results.sort(key=lambda x: x[1], reverse=True)
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Main hybrid search
+# ---------------------------------------------------------------------------
+
+
+def hybrid_search(
+    store: GraphStore,
+    query: str,
+    kind: Optional[str] = None,
+    limit: int = 20,
+    context_files: Optional[list[str]] = None,
+    model: Optional[str] = None,
+    provider: Optional[str] = None,
+) -> list[dict[str, Any]]:
+    """Hybrid search combining FTS5 BM25 and vector embeddings via RRF.
+
+    Attempts FTS5 + embedding search first, falling back to FTS5-only,
+    then keyword LIKE matching if FTS5 is unavailable.
+
+    Args:
+        store: The graph store to search.
+        query: Search query string.
+        kind: Optional node kind filter (e.g. ``"Function"``, ``"Class"``).
+        limit: Maximum results to return (default 20).
+        context_files: Optional list of file paths. Nodes in these files
+            receive a 1.5x score boost.
+
+    Returns:
+        List of dicts with node metadata and ``score`` field.
+    """
+    if not query or not query.strip():
+        return []
+
+    # NOTE: hybrid_search uses store._conn for FTS5 and keyword queries
+    # because those operate on the FTS virtual table or need raw Row
+    # access for batch-fetch performance.  This is documented coupling.
+    conn = store._conn
+    fetch_limit = limit * 3  # Fetch extra to allow for filtering and boosting
+
+    # ------ Phase 1: Gather ranked lists ------
+    fts_results: list[tuple[int, float]] = []
+    emb_results: list[tuple[int, float]] = []
+
+    # Try FTS5 search
+    try:
+        fts_results = _fts_search(conn, query, limit=fetch_limit)
+    except Exception as e:
+        logger.warning("FTS5 unavailable, will use fallback: %s", e)
+
+    # Try embedding search
+    emb_results = _embedding_search(
+        store, query, limit=fetch_limit, model=model, provider=provider,
+    )
+
+    # ------ Phase 2: Merge via RRF or fallback ------
+    if fts_results or emb_results:
+        lists_to_merge = []
+        if fts_results:
+            lists_to_merge.append(fts_results)
+        if emb_results:
+            lists_to_merge.append(emb_results)
+        merged = rrf_merge(*lists_to_merge)
+    else:
+        # Fallback: keyword LIKE matching
+        keyword_results = _keyword_search(conn, query, limit=fetch_limit)
+        if not keyword_results:
+            return []
+        merged = keyword_results
+
+    # ------ Phase 3+4: Batch-fetch nodes, apply boosting and kind filter ------
+    kind_boosts = detect_query_kind_boost(query)
+    context_set = set(context_files) if context_files else set()
+
+    # Batch-fetch all candidate nodes in one query
+    candidate_ids = [node_id for node_id, _ in merged]
+    node_rows: dict[int, Any] = {}
+    batch_size = 450
+    for i in range(0, len(candidate_ids), batch_size):
+        batch = candidate_ids[i:i + batch_size]
+        placeholders = ",".join("?" for _ in batch)
+        rows = conn.execute(
+            f"SELECT * FROM nodes WHERE id IN ({placeholders})",  # nosec B608
+            batch,
+        ).fetchall()
+        for row in rows:
+            node_rows[row["id"]] = row
+
+    # Apply boosting
+    boosted: list[tuple[int, float]] = []
+    for node_id, score in merged:
+        row = node_rows.get(node_id)
+        if not row:
+            continue
+
+        node_kind = row["kind"]
+        file_path = row["file_path"]
+        qualified_name = row["qualified_name"]
+
+        boost = 1.0
+        if node_kind in kind_boosts:
+            boost *= kind_boosts[node_kind]
+        if "_qualified" in kind_boosts and '.' in query:
+            if query.lower() in qualified_name.lower():
+                boost *= kind_boosts["_qualified"]
+        idents = kind_boosts.get("_qualified_identifiers")
+        if idents:
+            qn_lo = qualified_name.lower()
+            if any(ident in qn_lo for ident in idents):
+                boost *= 2.0
+        if context_set and file_path in context_set:
+            boost *= 1.5
+
+        boosted.append((node_id, score * boost))
+
+    boosted.sort(key=lambda x: x[1], reverse=True)
+
+    # Build results from the already-fetched rows
+    results: list[dict[str, Any]] = []
+    for node_id, final_score in boosted:
+        if len(results) >= limit:
+            break
+
+        row = node_rows.get(node_id)
+        if not row:
+            continue
+
+        node_kind = row["kind"]
+        if kind and node_kind != kind:
+            continue
+
+        results.append({
+            "name": _sanitize_name(row["name"]),
+            "qualified_name": _sanitize_name(row["qualified_name"]),
+            "kind": node_kind,
+            "file_path": row["file_path"],
+            "line_start": row["line_start"],
+            "line_end": row["line_end"],
+            "language": row["language"] or "",
+            "params": row["params"],
+            "return_type": row["return_type"],
+            "signature": row["signature"] if "signature" in row.keys() else None,
+            "score": round(final_score, 6),
+        })
+
+    return results