medgraph-engine/api/services/vector.py

"""Servicio de busqueda: semantica, full-text y hybrid con RRF."""

import os
from dotenv import load_dotenv
from services.graph import query

load_dotenv()

_client = None

# Pool size para retrieval inicial (se fusionan despues)
RETRIEVAL_POOL = 60
# Constante k para Reciprocal Rank Fusion
RRF_K = 60
# Embedding config
EMBEDDING_MODEL = "gemini-embedding-2-preview"
GCP_API_KEY = os.getenv("GCP_API_KEY", "")


def get_embedding_client():
    global _client
    if _client is None:
        from google import genai
        _client = genai.Client(api_key=GCP_API_KEY)
    return _client


def generate_embedding(text: str) -> list:
    client = get_embedding_client()
    truncated = text[:2000] if len(text) > 2000 else text
    r = client.models.embed_content(
        model=EMBEDDING_MODEL,
        contents=truncated,
    )
    return r.embeddings[0].values


def search_semantic(query_text: str, top_k: int = 40, libro_id: str = None) -> list:
    """Búsqueda semántica con vector KNN sobre embeddings."""
    embedding = generate_embedding(query_text)

    if libro_id:
        results = query("""
        CALL db.index.vector.queryNodes('chunk_embeddings', $top_k, $embedding)
        YIELD node AS c, score
        WHERE c.libro_id = $libro_id
        RETURN c.id AS id, c.libro_id AS libro, c.page_start AS pag_inicio,
               c.page_end AS pag_fin, c.text AS texto, c.word_count AS palabras,
               c.titulo_capitulo AS capitulo, c.titulo_seccion AS seccion,
               c.tipo_contenido AS tipo, c.parent_id AS parent_id, score
        ORDER BY score DESC
        """, {"embedding": embedding, "top_k": top_k, "libro_id": libro_id})
    else:
        results = query("""
        CALL db.index.vector.queryNodes('chunk_embeddings', $top_k, $embedding)
        YIELD node AS c, score
        RETURN c.id AS id, c.libro_id AS libro, c.page_start AS pag_inicio,
               c.page_end AS pag_fin, c.text AS texto, c.word_count AS palabras,
               c.titulo_capitulo AS capitulo, c.titulo_seccion AS seccion,
               c.tipo_contenido AS tipo, c.parent_id AS parent_id, score
        ORDER BY score DESC
        """, {"embedding": embedding, "top_k": top_k})

    return results


def search_keyword(query_text: str, top_k: int = 40, libro_id: str = None) -> list:
    """Búsqueda full-text con scoring BM25 sobre Chunk.text (Lucene via Neo4j)."""
    if not query_text.strip():
        return []

    # Lucene query: OR entre términos para mayor recall,
    # el scoring BM25 se encarga de rankear mejor los que tienen más matches
    terms = query_text.strip().split()
    lucene_query = " ".join(terms)  # OR implícito en Lucene

    params = {"query": lucene_query, "top_k": top_k}

    if libro_id:
        results = query("""
        CALL db.index.fulltext.queryNodes('busqueda_chunks', $query)
        YIELD node AS c, score
        WHERE c.libro_id = $libro_id
        RETURN c.id AS id, c.libro_id AS libro, c.page_start AS pag_inicio,
               c.page_end AS pag_fin, c.text AS texto, c.word_count AS palabras,
               c.titulo_capitulo AS capitulo, c.titulo_seccion AS seccion,
               c.tipo_contenido AS tipo, c.parent_id AS parent_id, score
        ORDER BY score DESC
        LIMIT $top_k
        """, params)
    else:
        results = query("""
        CALL db.index.fulltext.queryNodes('busqueda_chunks', $query)
        YIELD node AS c, score
        RETURN c.id AS id, c.libro_id AS libro, c.page_start AS pag_inicio,
               c.page_end AS pag_fin, c.text AS texto, c.word_count AS palabras,
               c.titulo_capitulo AS capitulo, c.titulo_seccion AS seccion,
               c.tipo_contenido AS tipo, c.parent_id AS parent_id, score
        ORDER BY score DESC
        LIMIT $top_k
        """, params)

    return results


def _rrf_fusion(rankings: list[list], k: int = RRF_K) -> list:
    """Reciprocal Rank Fusion: combina múltiples rankings en uno solo.

    Para cada resultado en cada ranking, calcula score = 1/(k + rank).
    Suma los scores de todos los rankings donde aparece cada chunk.
    Chunks que aparecen en múltiples rankings suben al top.
    """
    scores = {}
    chunk_data = {}

    for ranking in rankings:
        for rank, result in enumerate(ranking):
            chunk_id = result["id"]
            rrf_score = 1.0 / (k + rank + 1)

            if chunk_id not in scores:
                scores[chunk_id] = 0.0
                chunk_data[chunk_id] = result
            scores[chunk_id] += rrf_score

    # Ordenar por score RRF combinado
    sorted_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)

    results = []
    for chunk_id in sorted_ids:
        result = chunk_data[chunk_id]
        result["rrf_score"] = round(scores[chunk_id], 6)
        results.append(result)

    return results


def search_hybrid(query_text: str, top_k: int = 8, libro_id: str = None) -> dict:
    """Búsqueda híbrida con query rewriting + full-text + semántica + RRF.

    1. Query preprocessing (expansión sinónimos, clasificación intención, decomposición)
    2. Para cada sub-query: full-text top POOL + dense top POOL
    3. RRF fusion de todos los rankings
    4. Devuelve top_k resultados finales
    """
    from services.query import preprocess

    # 1. Preprocess query
    processed = preprocess(query_text)
    sub_queries = processed["sub_queries"]

    all_rankings = []
    total_keyword = 0
    total_semantic = 0

    # 2. Multi-retrieval por sub-query
    for sq in sub_queries:
        # Full-text
        try:
            kw = search_keyword(sq, top_k=RETRIEVAL_POOL, libro_id=libro_id)
            if kw:
                all_rankings.append(kw)
                total_keyword += len(kw)
        except Exception as e:
            import logging
            logging.error(f"Keyword search failed for '{sq}': {type(e).__name__}: {str(e)[:100]}")

        # Semantic (solo para la primera sub-query para no hacer muchas llamadas al embedding API)
        if sq == sub_queries[0]:
            try:
                sem = search_semantic(processed["expandida"], top_k=RETRIEVAL_POOL, libro_id=libro_id)
                if sem:
                    all_rankings.append(sem)
                    total_semantic += len(sem)
            except Exception as e:
                import logging
                logging.error(f"Semantic search failed: {type(e).__name__}: {str(e)[:100]}")

    if not all_rankings:
        return {
            "keyword_count": 0,
            "semantic_count": 0,
            "intencion": processed["intencion"],
            "query_expandida": processed["expandida"],
            "results": [],
        }

    # 3. RRF fusion de todos los rankings
    fused = _rrf_fusion(all_rankings)

    # 4. Top-k final
    results = fused[:top_k]

    return {
        "keyword_count": total_keyword,
        "semantic_count": total_semantic,
        "intencion": processed["intencion"],
        "query_expandida": processed["expandida"],
        "results": results,
    }