Hybrid Search + Reranking: do BM25 ao cross-encoder

# BM25 com PostgreSQL full-text search (ts_rank_cd é BM25-like)
import psycopg

def bm25_search(conn, query: str, k: int = 50) -> list[dict]:
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, chunk, ts_rank_cd(
              to_tsvector('portuguese', chunk),
              plainto_tsquery('portuguese', %s)
            ) AS score
            FROM chunks
            WHERE to_tsvector('portuguese', chunk)
                  @@ plainto_tsquery('portuguese', %s)
            ORDER BY score DESC
            LIMIT %s
        """, (query, query, k))
        return [{"id": r[0], "chunk": r[1], "score": r[2]} for r in cur.fetchall()]

# Para controle maior, use a extensão pg_bm25 (ParadeDB) ou
# um engine dedicado: Elasticsearch, OpenSearch, Tantivy, Meilisearch

# RRF: uma função de 10 linhas que bate esquemas elaborados
from collections import defaultdict

def rrf(rankings: list[list[str]], k: int = 60) -> list[tuple[str, float]]:
    """rankings: lista de listas de doc_ids ordenados por relevância."""
    scores: dict[str, float] = defaultdict(float)
    for ranking in rankings:
        for rank, doc_id in enumerate(ranking, start=1):
            scores[doc_id] += 1.0 / (k + rank)
    return sorted(scores.items(), key=lambda x: x[1], reverse=True)

# Uso
bm25_ids   = [r["id"] for r in bm25_search(conn, query, k=50)]
vector_ids = [r["id"] for r in vector_search(conn, query, k=50)]

fused = rrf([bm25_ids, vector_ids], k=60)
top_candidates = [doc_id for doc_id, _ in fused[:50]]

# Cohere rerank — simples, multilíngue, produção-ready
import cohere

co = cohere.ClientV2()

def rerank(query: str, docs: list[dict], top_n: int = 10) -> list[dict]:
    r = co.rerank(
        model="rerank-v3.5",
        query=query,
        documents=[d["chunk"] for d in docs],
        top_n=top_n,
    )
    # r.results: lista ordenada com índice do doc original + score
    return [
        {**docs[item.index], "rerank_score": item.relevance_score}
        for item in r.results
    ]

# Self-host com sentence-transformers
from sentence_transformers import CrossEncoder
ce = CrossEncoder("BAAI/bge-reranker-v2-m3", max_length=512)

def rerank_local(query: str, docs: list[dict], top_n: int = 10) -> list[dict]:
    pairs = [(query, d["chunk"]) for d in docs]
    scores = ce.predict(pairs, batch_size=32)
    ranked = sorted(zip(docs, scores), key=lambda x: x[1], reverse=True)
    return [{**d, "rerank_score": float(s)} for d, s in ranked[:top_n]]

# HyDE: Hypothetical Document Embedding
from anthropic import Anthropic

client = Anthropic()

def hyde_query(user_query: str) -> str:
    r = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=300,
        messages=[{
            "role": "user",
            "content": (
                "Escreva um parágrafo que responderia a esta pergunta como se "
                "fosse um trecho de documentação interna. 100-150 palavras, tom "
                "técnico e direto.\n\nPergunta: " + user_query
            ),
        }],
    )
    return r.content[0].text

# Pipeline com HyDE
def retrieve_with_hyde(query: str, k: int = 50) -> list[dict]:
    hypothetical_doc = hyde_query(query)        # LLM barato
    q_vec = embedder.encode(hypothetical_doc)    # embeda o doc, não a query
    return vector_search_by_vec(q_vec, k=k)

# Custo: +1 chamada Haiku (~US$0.0003 por query). Em alto volume,
# cacheie o HyDE por hash da query normalizada.

import numpy as np

def mmr(query_vec, doc_vecs, doc_ids, k: int = 5, lambda_: float = 0.7):
    """Maximal Marginal Relevance: rerank por relevância + diversidade."""
    sim_q = doc_vecs @ query_vec               # similaridade com query
    selected: list[int] = []
    candidates = list(range(len(doc_ids)))
    for _ in range(min(k, len(candidates))):
        if not selected:
            best = int(np.argmax(sim_q[candidates]))
        else:
            sim_sel = doc_vecs[candidates] @ doc_vecs[selected].T   # vs já escolhidos
            max_sim_sel = sim_sel.max(axis=1)
            score = lambda_ * sim_q[candidates] - (1 - lambda_) * max_sim_sel
            best = int(np.argmax(score))
        selected.append(candidates.pop(best))
    return [doc_ids[i] for i in selected]

# lambda=1 → recall puro (sem diversificar). lambda=0 → só diversidade (sem relevância).
# 0.6-0.8 é o range útil.

# Retrieval de 3 estágios: BM25 + vector → RRF → cross-encoder rerank
from collections import defaultdict

def retrieve(query: str, k_final: int = 8) -> list[dict]:
    # 1) Query transform (opcional — HyDE se query curta)
    search_query = query if len(query.split()) > 6 else hyde_query(query)

    # 2) Multi-method retrieve (paralelize com asyncio em prod)
    bm25_hits   = bm25_search(conn, search_query, k=50)
    vector_hits = vector_search(conn, search_query, k=50)
    by_id = {h["id"]: h for h in (*bm25_hits, *vector_hits)}

    # 3) RRF fusion
    fused = rrf(
        [[h["id"] for h in bm25_hits], [h["id"] for h in vector_hits]],
        k=60,
    )
    candidates = [by_id[doc_id] for doc_id, _ in fused[:50]]

    # 4) Cross-encoder rerank (top-50 → top-8)
    reranked = rerank(query, candidates, top_n=k_final)  # usa query original
    return reranked

-- pgvector com filtro ANTES da busca vetorial
SELECT id, chunk,
       1 - (embedding <=> $1::vector) AS sim
FROM chunks
WHERE tenant_id = $2
  AND lang = 'pt'
  AND doc_type = ANY($3::text[])
  AND created_at > now() - interval '90 days'
ORDER BY embedding <=> $1::vector
LIMIT 50;

-- Para HNSW pré-filtrar bem, use índice parcial ou partitioning.
-- Em Pinecone/Weaviate, metadata filter é campo estruturado no payload.