leocrm/app/plugins/builtins/unified_search/search_engine.py

"""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.config import settings
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"),
    "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_multi(
    result_lists: list[tuple[str, list[dict[str, Any]]]],
    k: int = 60,
) -> list[dict[str, Any]]:
    """Reciprocal Rank Fusion over N result lists.

    Each tuple is (mode_name, results). All lists get equal weight.
    score = sum(1/(k+rank_i) for each list where item appears)
    """
    fused: dict[str, dict[str, Any]] = {}

    for _mode_name, results in result_lists:
        for rank, item in enumerate(results):
            eid = str(item.get("id", ""))
            if not eid:
                continue
            rrf_score = 1.0 / (k + rank + 1)
            if eid not in fused:
                fused[eid] = {**item, "_score": 0.0}
            fused[eid]["_score"] += rrf_score

    return sorted(fused.values(), key=lambda x: x.get("_score", 0.0), reverse=True)


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.

    Backward-compatible wrapper around :func:`rrf_fusion_multi`.
    score = alpha * (1/(k+rank_fts)) + beta * (1/(k+rank_vec))
    """
    fused = rrf_fusion_multi(
        [("fts", fts_results), ("vec", vec_results)],
        k=k,
    )
    for item in fused:
        item["_entity_type"] = entity_type
    return fused


async def hybrid_search(
    db: AsyncSession,
    query_analysis: dict[str, Any],
    tenant_id: uuid.UUID,
    entity_types: list[str] | None = None,
    limit: int = 20,
    user_id: uuid.UUID | None = None,
    is_system_admin: bool = False,
) -> 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.

    Passes user_id and is_system_admin to each provider for visibility filtering.
    """
    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, db=db, tenant_id=tenant_id)

    from app.core.hooks import apply_filters
    query_analysis = await apply_filters("search.before_search", query_analysis)
    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]] = []
        rag_results: list[dict[str, Any]] = []
        graph_results: list[dict[str, Any]] = []

        # Respect provider capability flags
        if getattr(provider, "supports_fts", True):
            try:
                fts_results = await provider.search_fts(db, tsquery, tenant_id, fetch_limit, user_id=user_id, is_system_admin=is_system_admin)
            except Exception:
                logger.exception("FTS search failed for %s", entity_type)

        if query_embedding and getattr(provider, "supports_vector", True):
            try:
                vec_results = await provider.search_vector(db, query_embedding, tenant_id, fetch_limit, user_id=user_id, is_system_admin=is_system_admin)
            except Exception:
                logger.exception("Vector search failed for %s", entity_type)

        # RAG / Graph modes are not implemented yet, but keep the fusion ready.
        if getattr(provider, "supports_rag", False) and hasattr(provider, "search_rag") and query_embedding:
            try:
                rag_results = await provider.search_rag(db, query_embedding, tenant_id, fetch_limit, user_id=user_id, is_system_admin=is_system_admin)
            except Exception:
                logger.exception("RAG search failed for %s", entity_type)

        if getattr(provider, "supports_graph", False) and hasattr(provider, "search_graph"):
            try:
                graph_results = await provider.search_graph(db, query_analysis, tenant_id, fetch_limit, user_id=user_id, is_system_admin=is_system_admin)
            except Exception:
                logger.exception("Graph search failed for %s", entity_type)

        if rag_results or graph_results:
            fused = rrf_fusion_multi(
                [
                    ("fts", fts_results),
                    ("vec", vec_results),
                    ("rag", rag_results),
                    ("graph", graph_results),
                ]
            )
            for item in fused:
                item["_entity_type"] = entity_type
        else:
            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)
            # Preserve raw fields for post-search filtering (date/tags)
            result["_created_at"] = item.get("created_at")
            result["_updated_at"] = item.get("updated_at")
            result["_tags"] = item.get("tags")
            all_results.append(result)

    all_results.sort(key=lambda x: x.get("score", 0.0), reverse=True)
    from app.core.hooks import apply_filters
    all_results = await apply_filters("search.after_search", all_results)
    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"])

    # Set HNSW ef_search parameter for this transaction
    await db.execute(text(f"SET LOCAL hnsw.ef_search = {settings.hnsw_ef_search}"))

    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]