Búsqueda Híbrida con RRF Fusion

-- Migración: añadir soporte híbrido a tabla documents
-- LegalMind AI · PostgreSQL 16 + pgvector 0.7

CREATE EXTENSION IF NOT EXISTS vector;

-- Columna tsvector generada para búsqueda FTS
ALTER TABLE documents
ADD COLUMN IF NOT EXISTS ts_content tsvector
  GENERATED ALWAYS AS (
    to_tsvector('spanish',
      coalesce(titulo, '') || ' ' ||
      coalesce(contenido, '')
    )
  ) STORED;

-- Índice GIN para FTS (búsqueda de texto completo)
CREATE INDEX CONCURRENTLY IF NOT EXISTS
  idx_documents_fts
ON documents USING gin(ts_content);

-- Índice HNSW para vectorial (ya existente, verificar)
CREATE INDEX CONCURRENTLY IF NOT EXISTS
  idx_documents_embedding_hnsw
ON documents
USING hnsw(embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- Estadísticas para el planificador
ANALYZE documents;

import asyncpg
from typing import List, Dict, Optional

class LegalMindHybridSearch:
    """Búsqueda híbrida para corpus jurídico español."""

    def __init__(self, pool: asyncpg.Pool):
        self.pool = pool
        self.rrf_k = 60         # constante RRF
        self.v_weight = 0.6    # peso vectorial
        self.k_weight = 0.4    # peso BM25
        self.candidates = 30   # candidatos por canal

    async def search(
        self,
        query: str,
        embedding: List[float],
        top_k: int = 5,
        tipo_doc: Optional[str] = None
    ) -> List[Dict]:
        async with self.pool.acquire() as conn:
            rows = await conn.fetch("""
WITH vec AS (
  SELECT id, titulo, contenido,
    ROW_NUMBER() OVER (
      ORDER BY embedding <=> $1::vector
    ) AS v_rank
  FROM documents
  WHERE ($4 IS NULL OR tipo = $4)
  ORDER BY embedding <=> $1::vector
  LIMIT $3
),
fts AS (
  SELECT id, titulo, contenido,
    ROW_NUMBER() OVER (
      ORDER BY ts_rank(ts_content,
        websearch_to_tsquery('spanish', $2)
      ) DESC
    ) AS k_rank
  FROM documents
  WHERE ts_content @@
    websearch_to_tsquery('spanish', $2)
    AND ($4 IS NULL OR tipo = $4)
  ORDER BY ts_rank(ts_content,
    websearch_to_tsquery('spanish', $2)
  ) DESC
  LIMIT $3
)
SELECT
  COALESCE(v.id, f.id)           AS id,
  COALESCE(v.titulo, f.titulo)   AS titulo,
  COALESCE(v.contenido, f.contenido) AS contenido,
  -- RRF score ponderado
  COALESCE($5 / (60.0 + v.v_rank), 0) +
  COALESCE($6 / (60.0 + f.k_rank), 0) AS rrf_score
FROM vec v
FULL OUTER JOIN fts f ON v.id = f.id
ORDER BY rrf_score DESC
LIMIT $7
""",
                embedding, query,
                self.candidates, tipo_doc,
                self.v_weight, self.k_weight,
                top_k)
            return [dict(r) for r in rows]

from sentence_transformers import CrossEncoder
from functools import lru_cache
from typing import List, Dict

@lru_cache(maxsize=1)
def get_cross_encoder():
    # Cargado una sola vez en memoria
    return CrossEncoder(
        'cross-encoder/ms-marco-MiniLM-L-6-v2',
        max_length=512
    )

def rerank(
    query: str,
    candidates: List[Dict],
    top_k: int = 5
) -> List[Dict]:
    """
    Reordena candidatos con cross-encoder.
    Activar cuando rerank=True en el endpoint.
    """
    if not candidates:
        return candidates

    model = get_cross_encoder()
    pairs = [(query, c["contenido"][:400])
             for c in candidates]

    scores = model.predict(
        pairs,
        batch_size=16,
        show_progress_bar=False
    )

    for c, s in zip(candidates, scores):
        c["rerank_score"] = float(s)

    return sorted(
        candidates,
        key=lambda x: x["rerank_score"],
        reverse=True
    )[:top_k]

from fastapi import APIRouter, Depends
from pydantic import BaseModel
from openai import AsyncOpenAI

router = APIRouter()
openai = AsyncOpenAI()

class SearchRequest(BaseModel):
    query: str
    top_k: int = 5
    tipo_doc: Optional[str] = None
    rerank: bool = False

@router.post("/search")
async def search(
    req: SearchRequest,
    searcher: LegalMindHybridSearch = Depends(
        get_searcher
    )
):
    # 1. Embed la query
    resp = await openai.embeddings.create(
        model="text-embedding-3-small",
        input=req.query
    )
    emb = resp.data[0].embedding

    # 2. Búsqueda híbrida RRF
    results = await searcher.search(
        query=req.query,
        embedding=emb,
        top_k=req.top_k * 4 if req.rerank else req.top_k,
        tipo_doc=req.tipo_doc
    )

    # 3. Reranking opcional (premium)
    if req.rerank:
        results = rerank(
            req.query, results, req.top_k
        )

    return {
        "results": results,
        "method": "hybrid_rrf" +
                   ("+rerank" if req.rerank else "")
    }