"""Hybrid search engine: FTS + pgvector with RRF fusion.""" from __future__ import annotations import logging import uuid from typing import Any from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession from app.plugins.builtins.unified_search.embedding import generate_embedding from app.plugins.builtins.unified_search.provider_registry import get_search_registry logger = logging.getLogger(__name__) # Entity type -> (table_name, tsv_column, embedding_column) SEARCHABLE_ENTITIES: dict[str, tuple[str, str, str]] = { "contact": ("contacts", "search_tsv", "embedding"), "company": ("companies", "search_tsv", "embedding"), "mail": ("mails", "body_tsv", "embedding"), "file": ("files", "content_tsv", "embedding"), "event": ("calendar_entries", "search_tsv", "embedding"), } RRF_K = 60 RRF_ALPHA = 0.5 RRF_BETA = 0.5 def rrf_fusion( fts_results: list[dict[str, Any]], vec_results: list[dict[str, Any]], entity_type: str, k: int = RRF_K, alpha: float = RRF_ALPHA, beta: float = RRF_BETA, ) -> list[dict[str, Any]]: """Reciprocal Rank Fusion of FTS and vector search results. score = alpha * (1/(k+rank_fts)) + beta * (1/(k+rank_vec)) """ fused: dict[str, dict[str, Any]] = {} for rank, item in enumerate(fts_results): eid = str(item.get("id", "")) if not eid: continue rrf_score = alpha * (1.0 / (k + rank + 1)) if eid not in fused: fused[eid] = {**item, "_score": 0.0, "_entity_type": entity_type} fused[eid]["_score"] += rrf_score for rank, item in enumerate(vec_results): eid = str(item.get("id", "")) if not eid: continue rrf_score = beta * (1.0 / (k + rank + 1)) if eid not in fused: fused[eid] = {**item, "_score": 0.0, "_entity_type": entity_type} fused[eid]["_score"] += rrf_score return sorted(fused.values(), key=lambda x: x.get("_score", 0.0), reverse=True) async def hybrid_search( db: AsyncSession, query_analysis: dict[str, Any], tenant_id: uuid.UUID, entity_types: list[str] | None = None, limit: int = 20, ) -> list[dict[str, Any]]: """Perform hybrid search across all entity types. For each entity: FTS search + vector search + RRF fusion. Merge all results, sort by fused score, return top limit. """ registry = get_search_registry() all_entity_types = entity_types or registry.get_entity_types() normalized_query = query_analysis.get("normalized_query", "") semantic_terms = query_analysis.get("semantic_terms", []) # Build tsquery string tsquery_parts = [normalized_query] + semantic_terms tsquery = " & ".join( part.strip().replace(" ", " & ") for part in tsquery_parts if part and part.strip() ) if not tsquery: tsquery = normalized_query # Generate query embedding query_text = normalized_query if semantic_terms: query_text = f"{normalized_query} {' '.join(semantic_terms)}" query_embedding = await generate_embedding(query_text) all_results: list[dict[str, Any]] = [] fetch_limit = limit * 2 for entity_type in all_entity_types: provider = registry.get(entity_type) if provider is None: continue fts_results: list[dict[str, Any]] = [] vec_results: list[dict[str, Any]] = [] try: fts_results = await provider.search_fts(db, tsquery, tenant_id, fetch_limit) except Exception: logger.exception("FTS search failed for %s", entity_type) if query_embedding: try: vec_results = await provider.search_vector(db, query_embedding, tenant_id, fetch_limit) except Exception: logger.exception("Vector search failed for %s", entity_type) fused = rrf_fusion(fts_results, vec_results, entity_type) # Convert to search result format for item in fused: result = provider.to_search_result(item) result["score"] = item.get("_score", 0.0) all_results.append(result) all_results.sort(key=lambda x: x.get("score", 0.0), reverse=True) return all_results[:limit] async def find_similar_all_types( db: AsyncSession, entity_type: str, entity_id: uuid.UUID, tenant_id: uuid.UUID, limit: int = 5, ) -> dict[str, list[dict[str, Any]]]: """Find similar entities across all types based on embedding. Gets the embedding of the source entity, then searches all other tables. """ table_map = SEARCHABLE_ENTITIES source_info = table_map.get(entity_type) if not source_info: return {} table_name, _, emb_col = source_info # Get source embedding sql = text( f"SELECT {emb_col} AS embedding FROM {table_name} " f"WHERE id = :eid AND tenant_id = :tid" ) result = await db.execute(sql, {"eid": entity_id, "tid": tenant_id}) row = result.mappings().first() if not row or not row.get("embedding"): return {} source_embedding_str = str(row["embedding"]) registry = get_search_registry() similar: dict[str, list[dict[str, Any]]] = {} for etype, (tbl, _, emb) in table_map.items(): if etype == entity_type: continue provider = registry.get(etype) if provider is None: continue try: sql_sim = text( f""" SELECT *, 1 - ({emb} <=> cast(:emb_str AS vector)) AS score FROM {tbl} WHERE tenant_id = :tid AND deleted_at IS NULL AND {emb} IS NOT NULL ORDER BY {emb} <=> cast(:emb_str AS vector) LIMIT :lim """ ) res = await db.execute( sql_sim, {"emb_str": source_embedding_str, "tid": tenant_id, "lim": limit}, ) rows = res.mappings().all() results = [] for r in rows: sr = provider.to_search_result(dict(r)) sr["score"] = float(r.get("score", 0.0)) results.append(sr) similar[etype] = results except Exception: logger.exception("Similar search failed for %s", etype) similar[etype] = [] return similar async def autocomplete( db: AsyncSession, query: str, tenant_id: uuid.UUID, limit: int = 10, ) -> list[str]: """Autocomplete using FTS prefix search. Uses to_tsquery with :* suffix for prefix matching. """ query_clean = query.strip().replace(" ", " & ") if not query_clean: return [] tsquery = f"{query_clean}:*" suggestions: list[str] = [] registry = get_search_registry() for entity_type in registry.get_entity_types(): provider = registry.get(entity_type) if provider is None: continue entity_info = SEARCHABLE_ENTITIES.get(entity_type) if not entity_info: continue table_name, tsv_col, _ = entity_info try: sql = text( f""" SELECT * FROM {table_name} WHERE tenant_id = :tid AND deleted_at IS NULL AND {tsv_col} @@ to_tsquery('pg_catalog.german', :q) LIMIT :lim """ ) result = await db.execute( sql, {"q": tsquery, "tid": tenant_id, "lim": limit}, ) rows = result.mappings().all() for r in rows: sr = provider.to_search_result(dict(r)) title = sr.get("title", "") if title and title not in suggestions: suggestions.append(title) except Exception: logger.debug("Autocomplete failed for %s", entity_type) return suggestions[:limit]