Hybrid search + reranking: BM25 + dense + cross-encoder

from collections import defaultdict
from typing import Sequence

def rrf_fusion(rankings: Sequence[Sequence[str]], k: int = 60, top_n: int = 100):
    """
    rankings: lista de rankings, cada ranking é lista de doc_ids ordenada (rank 1 = posição 0)
    Retorna: lista de (doc_id, fused_score) ordenada.
    """
    scores: dict[str, float] = defaultdict(float)
    for ranking in rankings:
        for rank0, doc_id in enumerate(ranking):
            scores[doc_id] += 1.0 / (k + rank0 + 1)  # rank é 1-indexed
    return sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_n]

# Exemplo
bm25_ranking   = ["doc_42", "doc_3",  "doc_17", "doc_91", "doc_8"]   # top-5 BM25
vector_ranking = ["doc_3",  "doc_91", "doc_42", "doc_7",  "doc_17"]  # top-5 dense
fused = rrf_fusion([bm25_ranking, vector_ranking], k=60, top_n=5)

# Resultado:
# [('doc_3',  0.0326), ('doc_42', 0.0323),
#  ('doc_17', 0.0320), ('doc_91', 0.0322), ('doc_7',  0.0156)]
# doc_3 e doc_42 dominam porque aparecem alto em AMBOS os rankings

def weighted_rrf(rankings_with_weights, k=60, top_n=100):
    """
    rankings_with_weights: [(ranking, peso), ...]
    """
    scores = defaultdict(float)
    for ranking, w in rankings_with_weights:
        for rank0, doc_id in enumerate(ranking):
            scores[doc_id] += w / (k + rank0 + 1)
    return sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_n]

# Caso: domínio técnico, BM25 é mais valioso para SKUs e códigos
fused = weighted_rrf([
    (bm25_ranking,   0.7),
    (vector_ranking, 0.3),
], k=60, top_n=20)

# Pesos devem ser calibrados via golden set!

# Opção 1: Cohere Rerank (API gerenciada, multilingual)
import cohere
co = cohere.Client()  # COHERE_API_KEY no env

candidates = [{"id": d, "text": doc_texts[d]} for d in fused[:50]]
response = co.rerank(
    model="rerank-multilingual-v3.0",
    query=query,
    documents=[c["text"] for c in candidates],
    top_n=10,
)
final_ids = [candidates[r.index]["id"] for r in response.results]

# Opção 2: BGE-reranker local (GPU, sem custo por query)
from sentence_transformers import CrossEncoder

model = CrossEncoder("BAAI/bge-reranker-v2-m3", device="cuda")
pairs = [(query, doc_texts[d]) for d in fused[:50]]
scores = model.predict(pairs, batch_size=32)
ranked = sorted(zip(fused[:50], scores), key=lambda x: x[1], reverse=True)
final_ids = [d for d, _ in ranked[:10]]

# Opção 3: Qwen reranker (multilingual, ótimo PT-BR em 2026)
from FlagEmbedding import FlagReranker
reranker = FlagReranker("Qwen/Qwen2.5-Reranker-0.5B", use_fp16=True)
scores = reranker.compute_score([[query, doc_texts[d]] for d in fused[:50]])

from elasticsearch import Elasticsearch
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer, CrossEncoder
from collections import defaultdict

# 1. Inicialização
es = Elasticsearch(["https://es.internal:9200"])
qdrant = QdrantClient(url="https://qdrant.internal")
embedder = SentenceTransformer("BAAI/bge-m3")
reranker = CrossEncoder("BAAI/bge-reranker-v2-m3", device="cuda")

def hybrid_search(query: str, top_k: int = 10, retrieval_k: int = 100, rerank_k: int = 50):
    # 2a. BM25 via Elasticsearch
    es_resp = es.search(
        index="docs",
        body={"query": {"match": {"content": query}}, "size": retrieval_k},
        timeout="2s",
    )
    bm25_ids = [hit["_id"] for hit in es_resp["hits"]["hits"]]

    # 2b. Dense via Qdrant
    q_emb = embedder.encode(query, normalize_embeddings=True).tolist()
    dense_resp = qdrant.search(
        collection_name="docs",
        query_vector=q_emb,
        limit=retrieval_k,
    )
    dense_ids = [hit.id for hit in dense_resp]

    # 3. RRF fusion
    fused_scores: dict[str, float] = defaultdict(float)
    for ranking in [bm25_ids, dense_ids]:
        for rank0, doc_id in enumerate(ranking):
            fused_scores[doc_id] += 1.0 / (60 + rank0 + 1)
    fused = [d for d, _ in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)][:rerank_k]

    # 4. Cross-encoder rerank dos top-K fundidos
    doc_texts = fetch_doc_texts(fused)  # batch fetch do storage
    pairs = [(query, doc_texts[d]) for d in fused]
    rerank_scores = reranker.predict(pairs, batch_size=32, show_progress_bar=False)
    final = sorted(zip(fused, rerank_scores), key=lambda x: x[1], reverse=True)[:top_k]

    return [{"id": d, "score": float(s)} for d, s in final]


# Latência típica em produção:
# - ES BM25:          5-15 ms
# - Qdrant ANN:       5-20 ms  (em paralelo com ES, mesmo tempo total)
# - RRF (in-memory):  <1 ms
# - Reranker (GPU):  100-200 ms (top-50 em batch)
# - Total:           ~120-230 ms p95