codegraph-nav 0.1.0__py3-none-any.whl

codegraph_nav/graph/search.py

@@ -0,0 +1,257 @@
+"""Hybrid search — FTS5 + fuzzy + RRF fusion.
+
+Combines four ranked lists via Reciprocal Rank Fusion (RRF):
+  A: Exact name match (score=1.0)
+  B: Fuzzy name match (existing CodeSearcher)
+  C: FTS5 BM25 over name/qualified_name/file_path/signature
+  D: (Reserved for Phase C concept/tag match)
+
+Gracefully degrades if FTS5 is not compiled into SQLite.
+"""
+
+from __future__ import annotations
+
+import re
+import sqlite3
+
+from .store import GraphStore
+
+# ==============================================================================
+# FTS5 Index Management
+# ==============================================================================
+
+
+def rebuild_fts_index(conn: sqlite3.Connection) -> int:
+    """Rebuild FTS5 full-text index from nodes table.
+
+    Returns number of indexed rows, or -1 if FTS5 unavailable.
+    """
+    try:
+        conn.execute("DROP TABLE IF EXISTS nodes_fts")
+        conn.execute("""CREATE VIRTUAL TABLE nodes_fts USING fts5(
+                name, qualified_name, file_path, signature,
+                tokenize='porter unicode61'
+            )""")
+        conn.execute("""INSERT INTO nodes_fts(rowid, name, qualified_name, file_path, signature)
+               SELECT id, name, qualified_name, file_path, COALESCE(signature, '')
+               FROM nodes""")
+        conn.commit()
+        count: int = conn.execute("SELECT COUNT(*) FROM nodes_fts").fetchone()[0]
+        return count
+    except sqlite3.OperationalError:
+        return -1
+
+
+# ==============================================================================
+# FTS5 Search
+# ==============================================================================
+
+
+def fts_search(
+    conn: sqlite3.Connection,
+    query: str,
+    limit: int = 50,
+) -> list[tuple[int, float]]:
+    """BM25 search via FTS5. Returns [(node_id, score)].
+
+    Query is wrapped in quotes to prevent FTS5 operator injection.
+    """
+    # Sanitize: wrap in quotes, escape internal quotes
+    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 (lower = better); negate for sorting
+        return [(row[0], -row[1]) for row in rows]
+    except sqlite3.OperationalError:
+        return []
+
+
+# ==============================================================================
+# Reciprocal Rank Fusion (RRF)
+# ==============================================================================
+
+
+def rrf_merge(
+    *result_lists: list[tuple[int, float]],
+    k: int = 60,
+) -> list[tuple[int, float]]:
+    """Merge multiple ranked lists via Reciprocal Rank Fusion.
+
+    RRF score for an item = sum(1 / (k + rank + 1)) across all lists
+    where 'rank' is 0-indexed position in each list.
+
+    Args:
+        *result_lists: Each list is [(item_id, score)], ordered by score desc.
+        k: Smoothing constant (default 60, higher = rank differences matter less).
+
+    Returns:
+        Merged list of (item_id, rrf_score), sorted by rrf_score desc.
+    """
+    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
+
+
+# ==============================================================================
+# Query Kind Boosting
+# ==============================================================================
+
+
+def detect_query_boosts(query: str) -> dict[str, float]:
+    """Detect query patterns and return kind boosts.
+
+    - PascalCase → boost Class 1.5x
+    - snake_case → boost Function 1.5x
+    - Contains "/" → boost file path matches 2.0x
+    """
+    boosts: dict[str, float] = {}
+    q = query.strip()
+
+    # PascalCase: starts with upper, has lower
+    if re.match(r"^[A-Z][a-z]", q) and not q.isupper():
+        boosts["Class"] = 1.5
+        boosts["Method"] = 1.2
+
+    # snake_case: has underscore and letters
+    if "_" in q and re.search(r"[a-zA-Z]", q):
+        boosts["Function"] = 1.5
+
+    # File path pattern
+    if "/" in q or "." in q:
+        boosts["_file_path"] = 2.0
+
+    return boosts
+
+
+# ==============================================================================
+# Hybrid Search
+# ==============================================================================
+
+
+def hybrid_search(
+    store: GraphStore,
+    query: str,
+    limit: int = 20,
+    k: int = 60,
+) -> list[dict]:
+    """Hybrid search combining FTS5 and name matching via RRF.
+
+    Lists fused:
+      A: Exact name match (from DB)
+      B: Prefix/contains name match (from DB)
+      C: FTS5 BM25 (if available)
+
+    Returns list of dicts: {id, name, qualified_name, file_path, kind, lines, score}.
+    """
+    conn = store.conn
+    query_lower = query.lower().strip()
+    boosts = detect_query_boosts(query)
+
+    # List A: Exact name match
+    exact_rows = conn.execute(
+        "SELECT id, name, kind FROM nodes WHERE LOWER(name) = ? LIMIT ?",
+        (query_lower, limit),
+    ).fetchall()
+    list_a = [(row[0], 1.0) for row in exact_rows]
+
+    # List B: Contains/prefix match
+    like_pattern = f"%{query_lower}%"
+    contains_rows = conn.execute(
+        "SELECT id, name, kind FROM nodes WHERE LOWER(name) LIKE ? " "AND LOWER(name) != ? LIMIT ?",
+        (like_pattern, query_lower, limit * 2),
+    ).fetchall()
+    list_b = []
+    for row in contains_rows:
+        name_lower = row[1].lower()
+        if name_lower.startswith(query_lower):
+            score = 0.8
+        elif query_lower in name_lower:
+            score = 0.5 + len(query_lower) / len(name_lower) * 0.3
+        else:
+            score = 0.3
+        list_b.append((row[0], score))
+    list_b.sort(key=lambda x: x[1], reverse=True)
+
+    # List C: FTS5 (if available)
+    list_c = []
+    if store.fts_available:
+        list_c = fts_search(conn, query, limit=limit * 2)
+
+    # Apply kind boosts
+    if boosts:
+        for lst in [list_a, list_b, list_c]:
+            boosted = []
+            for node_id, score in lst:
+                row = conn.execute(
+                    "SELECT kind, file_path FROM nodes WHERE id = ?", (node_id,)
+                ).fetchone()
+                if row:
+                    kind = row[0]
+                    boost = boosts.get(kind, 1.0)
+                    if "_file_path" in boosts and query_lower in (row[1] or "").lower():
+                        boost *= boosts["_file_path"]
+                    boosted.append((node_id, score * boost))
+                else:
+                    boosted.append((node_id, score))
+            lst.clear()
+            lst.extend(boosted)
+            lst.sort(key=lambda x: x[1], reverse=True)
+
+    # RRF merge
+    lists_to_merge = [l for l in [list_a, list_b, list_c] if l]
+    if not lists_to_merge:
+        return []
+
+    merged = rrf_merge(*lists_to_merge, k=k)
+
+    # Fetch full node info for top results
+    results = []
+    for node_id, rrf_score in merged[:limit]:
+        node = store.get_node_by_id(node_id)
+        if node:
+            results.append(
+                {
+                    "id": node["id"],
+                    "name": node["name"],
+                    "qualified_name": node["qualified_name"],
+                    "file_path": node["file_path"],
+                    "kind": node["kind"],
+                    "line_start": node["line_start"],
+                    "line_end": node["line_end"],
+                    "signature": node["signature"],
+                    "score": round(rrf_score, 6),
+                }
+            )
+
+    return results
+
+
+def format_search_results_minimal(results: list[dict], limit: int = 20) -> str:
+    """Format hybrid search results in compact format."""
+    if not results:
+        return "No matching symbols found."
+
+    type_abbr = {
+        "Function": "fn",
+        "Class": "cls",
+        "Method": "mth",
+        "Variable": "var",
+        "File": "file",
+    }
+
+    lines = [f"Found {len(results)} matches:"]
+    for r in results[:limit]:
+        abbr = type_abbr.get(r["kind"], r["kind"][:3])
+        end = r["line_end"] or r["line_start"] or "?"
+        lines.append(f"{r['file_path']}:L{r['line_start']}-{end} [{abbr}] {r['name']}")
+
+    if len(results) > limit:
+        lines.append(f"... +{len(results) - limit} more")
+
+    return "\n".join(lines)