wehub-resource-sync
448 lines
15 KiB
Python
448 lines
15 KiB
Python
"""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
|