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.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]