Hybrid Search

from rank_bm25 import BM25Okapi
import numpy as np

documents = [
    "OAuth2 authentication failure troubleshooting guide",
    "How to configure SSO with SAML providers",
    "Debugging login issues with identity providers",
    "REST API authentication best practices",
    "Token refresh flow implementation guide",
    "Kubernetes pod authentication with service accounts",
    "CORS preflight request failures in browser",
]

# Tokenize (production: use proper tokenizer with stemming)
tokenized_docs = [doc.lower().split() for doc in documents]

# Build BM25 index — no training data needed
bm25 = BM25Okapi(tokenized_docs, k1=1.5, b=0.75)

# Search
query = "authentication failure OAuth2"
tokenized_query = query.lower().split()
scores = bm25.get_scores(tokenized_query)

# Ranked results
ranked = np.argsort(scores)[::-1]
print("BM25 Results:")
for idx in ranked[:3]:
    print(f"  score={scores[idx]:.3f}: {documents[idx]}")

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer('BAAI/bge-small-en-v1.5')

documents = [
    "OAuth2 authentication failure troubleshooting guide",
    "How to configure SSO with SAML providers",
    "Debugging login issues with identity providers",
    "REST API authentication best practices",
    "Token refresh flow implementation guide",
    "Kubernetes pod authentication with service accounts",
    "CORS preflight request failures in browser",
]

# Encode all documents once (store these vectors)
doc_embeddings = model.encode(documents, normalize_embeddings=True)

# Encode query at search time
query = "fix login problems"
query_embedding = model.encode(query, normalize_embeddings=True)

# Cosine similarity = dot product (since normalized)
scores = np.dot(doc_embeddings, query_embedding)

ranked = np.argsort(scores)[::-1]
print("Vector Search Results:")
for idx in ranked[:3]:
    print(f"  score={scores[idx]:.3f}: {documents[idx]}")

from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
import numpy as np

def reciprocal_rank_fusion(
    ranked_lists: list[list[int]],
    k: int = 60
) -> list[tuple[int, float]]:
    """
    Combine multiple ranked lists using RRF.

    Args:
        ranked_lists: List of ranked document ID lists.
                      Each inner list is ordered by relevance (best first).
        k: RRF constant (default 60, from original paper).

    Returns:
        List of (doc_id, rrf_score) tuples, sorted by score descending.
    """
    rrf_scores: dict[int, float] = {}

    for ranked_list in ranked_lists:
        for rank, doc_id in enumerate(ranked_list, start=1):
            if doc_id not in rrf_scores:
                rrf_scores[doc_id] = 0.0
            rrf_scores[doc_id] += 1.0 / (k + rank)

    return sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)


def hybrid_search(
    query: str,
    documents: list[str],
    bm25: BM25Okapi,
    model: SentenceTransformer,
    doc_embeddings: np.ndarray,
    top_k: int = 20,
    k: int = 60
) -> list[tuple[int, float, float, float]]:
    """
    Full hybrid search pipeline.

    Returns: [(doc_idx, rrf_score, bm25_score, vector_score), ...]
    """
    # Path 1: BM25
    tokenized_query = query.lower().split()
    bm25_scores = bm25.get_scores(tokenized_query)
    bm25_ranking = np.argsort(-bm25_scores)[:top_k].tolist()

    # Path 2: Vector
    query_emb = model.encode(query, normalize_embeddings=True)
    vector_scores = np.dot(doc_embeddings, query_emb)
    vector_ranking = np.argsort(-vector_scores)[:top_k].tolist()

    # Fusion
    fused = reciprocal_rank_fusion([bm25_ranking, vector_ranking], k=k)

    # Attach original scores for debugging
    results = []
    for doc_idx, rrf_score in fused:
        results.append((
            doc_idx, rrf_score,
            float(bm25_scores[doc_idx]),
            float(vector_scores[doc_idx])
        ))

    return results


# --- Usage ---
documents = [
    "OAuth2 authentication failure troubleshooting guide",
    "How to configure SSO with SAML providers",
    "Debugging login issues with identity providers",
    "REST API authentication best practices",
    "Token refresh flow implementation guide",
]

tokenized_docs = [doc.lower().split() for doc in documents]
bm25 = BM25Okapi(tokenized_docs)
model = SentenceTransformer('BAAI/bge-small-en-v1.5')
doc_embeddings = model.encode(documents, normalize_embeddings=True)

query = "how to fix authentication failure in OAuth2"
results = hybrid_search(query, documents, bm25, model, doc_embeddings)

print("Hybrid Search Results:")
for idx, rrf, bm25_s, vec_s in results[:5]:
    print(f"  RRF={rrf:.5f} | BM25={bm25_s:.3f} | Vec={vec_s:.3f}")
    print(f"    {documents[idx]}")

def weighted_rrf(
    bm25_ranking: list[int],
    vector_ranking: list[int],
    alpha: float = 0.5,
    k: int = 60
) -> list[tuple[int, float]]:
    """
    Weighted RRF: alpha controls the balance.

    alpha = 0.5 -> equal weight (standard RRF)
    alpha = 0.7 -> 70% keyword, 30% semantic
    alpha = 0.3 -> 30% keyword, 70% semantic
    """
    scores: dict[int, float] = {}

    for rank, doc_id in enumerate(bm25_ranking, start=1):
        scores[doc_id] = scores.get(doc_id, 0.0) + alpha / (k + rank)

    for rank, doc_id in enumerate(vector_ranking, start=1):
        scores[doc_id] = scores.get(doc_id, 0.0) + (1 - alpha) / (k + rank)

    return sorted(scores.items(), key=lambda x: x[1], reverse=True)